Auto-sync: 2026-06-25 22:43:25 (part 17)

outputs/audit_venv/lib/python3.11/site-packages/torch/_inductor/memory.py

@@ -0,0 +1,1108 @@
+from __future__ import annotations
+
+import collections
+import dataclasses
+import heapq
+import logging
+from typing import TYPE_CHECKING, TypedDict
+
+import torch
+from torch._environment import is_fbcode
+from torch._utils_internal import signpost_event
+from torch.utils._ordered_set import OrderedSet
+
+from . import config
+from .ir import MultiOutputLayout, NoneLayout
+from .utils import get_dtype_size, is_nonfreeable_buffers
+from .virtualized import V
+
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from .dependencies import Dep
+    from .scheduler import BaseSchedulerNode, SchedulerBuffer
+
+from .dependencies import WeakDep
+
+
+torch_log = logging.getLogger(__name__)
+
+
+@dataclasses.dataclass
+class PeakMemoryResult:
+    order: list[BaseSchedulerNode]
+    peak_memory: int
+    method: str
+
+
+@dataclasses.dataclass
+class MemoryPlanningInfoForBuffer:
+    size_alloc: int = 0
+    size_free: int = 0
+    # succ_nodes used for buffer lifetime/freeing (excludes is_fake WeakDeps)
+    succ_nodes: OrderedSet[BaseSchedulerNode] = dataclasses.field(
+        default_factory=OrderedSet
+    )
+    # succ_nodes used for node ordering (includes is_fake WeakDeps)
+    succ_nodes_for_ordering: OrderedSet[BaseSchedulerNode] = dataclasses.field(
+        default_factory=OrderedSet
+    )
+
+    def __post_init__(self) -> None:
+        torch._check(
+            len(self.succ_nodes) <= len(self.succ_nodes_for_ordering),
+            lambda: f"succ_nodes must be a subset of succ_nodes_for_ordering. "
+            f"len(succ_nodes)={len(self.succ_nodes)}, len(succ_nodes_for_ordering)={len(self.succ_nodes_for_ordering)}",
+        )
+
+
+@dataclasses.dataclass
+class MemoryPlanningInfoForNode:
+    index: int = 0
+    size: int = 0
+    pred_buffers: OrderedSet[SchedulerBuffer | FreeableInputBuffer] = dataclasses.field(
+        default_factory=OrderedSet
+    )
+    pred_nodes: OrderedSet[BaseSchedulerNode] = dataclasses.field(
+        default_factory=OrderedSet
+    )
+    succ_nodes: OrderedSet[BaseSchedulerNode] = dataclasses.field(
+        default_factory=OrderedSet
+    )
+
+
+@dataclasses.dataclass
+class FreeableInputBuffer:
+    name: str
+    mpi_buffer: MemoryPlanningInfoForBuffer = dataclasses.field(
+        default_factory=MemoryPlanningInfoForBuffer
+    )
+
+    def get_name(self) -> str:
+        return self.name
+
+    def __hash__(self) -> int:
+        return hash(self.name)
+
+
+def get_freeable_input_buf(
+    nodes: list[BaseSchedulerNode],
+    graph_inputs: OrderedSet[str],
+) -> dict[str, FreeableInputBuffer]:
+    """
+    Create and keep track of all input buffers that can be freed during the program
+
+    Returns:
+        A dictionary containing all freeable input buffers, keyed by their names.
+    """
+
+    def _dep_size_hint(dep: Dep) -> int:
+        return V.graph.get_dep_size_hint(dep)
+
+    # get freeable input buffers' successor nodes for memory lifetime (excludes is_fake WeakDeps)
+    # and for ordering (includes all deps)
+    dep_name_to_succ_nodes: dict[str, OrderedSet[BaseSchedulerNode]] = (
+        collections.defaultdict(OrderedSet)
+    )
+    dep_name_to_succ_nodes_for_ordering: dict[str, OrderedSet[BaseSchedulerNode]] = (
+        collections.defaultdict(OrderedSet)
+    )
+    dep_name_to_size: dict[str, int] = dict()
+
+    for node in nodes:
+        for dep in node.read_writes.reads:
+            if dep.name in graph_inputs:
+                if not is_nonfreeable_buffers(dep):
+                    # All deps contribute to ordering, but fake weak deps do not contribute to
+                    # memory liveness
+                    dep_name_to_succ_nodes_for_ordering[dep.name].add(node)
+                    dep_name_to_size[dep.name] = _dep_size_hint(dep)
+                    if not (isinstance(dep, WeakDep) and dep.is_fake):
+                        dep_name_to_succ_nodes[dep.name].add(node)
+
+    # create FreeableInputBuffer objects and add them to the returned dictionary
+    name_to_freeable_input_buf: dict[str, FreeableInputBuffer] = dict()
+    for dep_name in dep_name_to_succ_nodes_for_ordering:
+        name_to_freeable_input_buf[dep_name] = FreeableInputBuffer(
+            dep_name,
+            MemoryPlanningInfoForBuffer(
+                size_free=dep_name_to_size[dep_name],
+                succ_nodes=dep_name_to_succ_nodes[dep_name],
+                succ_nodes_for_ordering=dep_name_to_succ_nodes_for_ordering[dep_name],
+            ),
+        )
+    return name_to_freeable_input_buf
+
+
+def compute_size_for_scheduler_buffer(
+    name_to_buf: dict[str, SchedulerBuffer],
+) -> dict[str, tuple[int, int]]:
+    """
+    Compute the size of each scheduler buffer, including (1) memory allocated when
+    it is created and (2) memory deallocated when it is freed.
+
+    We specially handle the case of MultiOutputLayout.
+    Consider the following case:
+        buf0 = some_ops_with_multi_outputs(...)
+        buf1 = buf0[0] # assume 10 bytes
+        buf2 = buf0[1] # assume 20 bytes
+    In such cases,
+        buf0: at creation, 30 bytes allocated, when deleted, 0 bytes freed
+        buf1: at creation, 0 bytes allocated, when deleted, 10 bytes freed
+        buf2: at creation, 0 bytes allocated, when deleted, 20 bytes freed
+
+    When an operation mutates a buffer in-place, the scheduler creates a new buffer name
+    to track the "before" and "after" states, even though they share the same memory.
+
+    The mutated buffer represents a rename with zero allocation and deallocation cost.
+    During dependency tracking, we transfer dependencies from the mutated name back to
+    the original buffer, ensuring the original memory is only freed when all aliases
+    are done.
+
+    This handles cases where a buffer has multiple non-overlapping aliases - rather than
+    trying to assign free costs to individual aliases, we forward all alias dependencies
+    to the original buffer.
+
+    Consider:
+        buf0 = op0()
+        buf1 = mutation_op_(buf0)
+        del buf0
+        ...
+        op(buf1)
+        del buf1
+
+    The only memory events are the creation prior to op0, and the deletion following buf1.
+
+    Returns:
+        A dictionary mapping a scheduler buffer to a tuple of (size_alloc, size_free).
+    """
+    from .ir import MultiOutput
+    from .scheduler import OutputNode
+
+    sched_buf_to_size: dict[str, tuple[int, int]] = dict()
+
+    def _compute_and_update_buf_size(
+        sched_buf: SchedulerBuffer, user_of_MultiOutputLayout: bool = False
+    ) -> int:
+        if sched_buf.get_name() in V.graph.scheduler.mutation_real_name:
+            sched_buf_to_size[sched_buf.get_name()] = (0, 0)
+            return 0
+        elif isinstance(sched_buf.node.layout, NoneLayout):
+            sched_buf_to_size[sched_buf.get_name()] = (0, 0)
+            return 0
+        elif isinstance(sched_buf.node.layout, MultiOutputLayout):
+            size_alloc = 0
+            for user in sched_buf.users:
+                if isinstance(user.node, OutputNode):
+                    continue
+                for buf in user.node.get_outputs():
+                    if isinstance(buf.node, MultiOutput):
+                        size_alloc += _compute_and_update_buf_size(buf, True)
+            sched_buf_to_size[sched_buf.get_name()] = (
+                0 if user_of_MultiOutputLayout else size_alloc,
+                0,
+            )
+            return size_alloc
+        else:
+            buf_size = V.graph.sizevars.optimization_hint(
+                sched_buf.node.get_numel(), fallback=0
+            ) * get_dtype_size(sched_buf.node.get_dtype())
+            sched_buf_to_size[sched_buf.get_name()] = (
+                0 if user_of_MultiOutputLayout else buf_size,
+                buf_size,
+            )
+            return buf_size
+
+    for sched_buf in name_to_buf.values():
+        # skip if sched_buf is already processed as an user of another SchedulerBuffer
+        # whose layout is of the type MultiOutputLayout
+        if sched_buf.get_name() not in sched_buf_to_size:
+            _compute_and_update_buf_size(sched_buf)
+
+    return sched_buf_to_size
+
+
+def assign_memory_planning_info_for_scheduler_buffers(
+    nodes: list[BaseSchedulerNode],
+    name_to_buf: dict[str, SchedulerBuffer],
+) -> None:
+    """
+    For each SchedulerBuffer, assign its size info and successor nodes.
+    A buffer's successor nodes determines when a buffer can be freed.
+    """
+    # get buffer sizes
+    sched_buf_to_size = compute_size_for_scheduler_buffer(name_to_buf)
+
+    # get buffer's successor nodes for memory lifetime (excludes is_fake WeakDeps)
+    # and for ordering (includes all deps)
+    dep_name_to_succ_nodes: dict[str, OrderedSet[BaseSchedulerNode]] = (
+        collections.defaultdict(OrderedSet)
+    )
+    dep_name_to_succ_nodes_for_ordering: dict[str, OrderedSet[BaseSchedulerNode]] = (
+        collections.defaultdict(OrderedSet)
+    )
+    for node in nodes:
+        for dep in node.unmet_dependencies:
+            # All deps contribute to ordering, but fake weak deps do not contribute to
+            # memory liveness
+            dep_name_to_succ_nodes_for_ordering[dep.name].add(node)
+            if not (isinstance(dep, WeakDep) and dep.is_fake):
+                dep_name_to_succ_nodes[dep.name].add(node)
+
+    # iterate in reverse, so dependencies are picked up transitively.
+    for mutating_buf_name, real_buf_name in reversed(
+        V.graph.scheduler.mutation_real_name.items()
+    ):
+        dep_name_to_succ_nodes[real_buf_name] |= dep_name_to_succ_nodes[
+            mutating_buf_name
+        ]
+        dep_name_to_succ_nodes_for_ordering[real_buf_name] |= (
+            dep_name_to_succ_nodes_for_ordering[mutating_buf_name]
+        )
+
+    # populate the MemoryPlanningInfoForBuffer attribute to each scheduler buffer
+    # note: there are scheduler buffers not in dep_name_to_succ_nodes (e.g., graph outputs)
+    for buf_name in name_to_buf:
+        name_to_buf[buf_name].mpi_buffer = MemoryPlanningInfoForBuffer(
+            size_alloc=sched_buf_to_size[buf_name][0],
+            size_free=sched_buf_to_size[buf_name][1],
+            succ_nodes=dep_name_to_succ_nodes[buf_name],
+            succ_nodes_for_ordering=dep_name_to_succ_nodes_for_ordering[buf_name],
+        )
+
+
+def assign_memory_planning_info_for_scheduler_nodes(
+    nodes: list[BaseSchedulerNode],
+    name_to_fused_node: dict[str, BaseSchedulerNode],
+    name_to_buf: dict[str, SchedulerBuffer],
+    name_to_freeable_input_buf: dict[str, FreeableInputBuffer],
+) -> None:
+    """
+    Assign to each scheduler node its predecessor and successor nodes.
+    """
+
+    node_to_pred_nodes: dict[BaseSchedulerNode, OrderedSet[BaseSchedulerNode]] = (
+        collections.defaultdict(OrderedSet)
+    )
+    node_to_succ_nodes: dict[BaseSchedulerNode, OrderedSet[BaseSchedulerNode]] = {}
+    node_to_pred_buffers: dict[
+        BaseSchedulerNode, OrderedSet[SchedulerBuffer | FreeableInputBuffer]
+    ] = collections.defaultdict(OrderedSet)
+
+    # collect all predecessors using existing successor mappings
+    for node in nodes:
+        succ_nodes = OrderedSet(
+            succ_node
+            for buffer in node.get_outputs()
+            for succ_node in buffer.mpi_buffer.succ_nodes_for_ordering
+        )
+        node_to_succ_nodes[node] = succ_nodes
+
+        # For each successor, add current node as its predecessor
+        for succ_node in succ_nodes:
+            node_to_pred_nodes[succ_node].add(node)
+
+        # For each output buffer, add it as predecessor to its successor nodes
+        # Use succ_nodes (not succ_nodes_for_ordering) since pred_buffers is used
+        # for memory lifetime tracking, not ordering
+        for buffer in node.get_outputs():
+            for succ_node in buffer.mpi_buffer.succ_nodes:
+                node_to_pred_buffers[succ_node].add(buffer)
+
+    for freeable_buffer in name_to_freeable_input_buf.values():
+        for succ_node in freeable_buffer.mpi_buffer.succ_nodes:
+            node_to_pred_buffers[succ_node].add(freeable_buffer)
+
+    # Second pass: assign memory planning info using completed predecessor mappings
+    for index, node in enumerate(nodes):
+        size_alloc = sum(buffer.mpi_buffer.size_alloc for buffer in node.get_outputs())
+        succ_nodes = node_to_succ_nodes[node]
+        pred_nodes = node_to_pred_nodes[node]
+
+        # make sure we do not make node a successor or predecessor of itself
+        succ_nodes.discard(node)
+        pred_nodes.discard(node)
+
+        node.mpi_node = MemoryPlanningInfoForNode(
+            index=index,
+            size=size_alloc,
+            pred_buffers=node_to_pred_buffers[node],
+            pred_nodes=node_to_pred_nodes[node],
+            succ_nodes=succ_nodes,
+        )
+
+
+# map each scheduler buffer to its size, start step, and end step
+@dataclasses.dataclass
+class BufferInfo:
+    buffer: SchedulerBuffer | FreeableInputBuffer
+    size_alloc: int
+    size_free: int
+    start_step: int
+    end_step: int
+
+
+def compute_memory_timeline(
+    nodes: list[BaseSchedulerNode],
+    name_to_freeable_input_buf: dict[str, FreeableInputBuffer],
+    graph_outputs: OrderedSet[str],
+) -> tuple[
+    list[BufferInfo],
+    dict[BaseSchedulerNode, int],
+    dict[FreeableInputBuffer | SchedulerBuffer, BaseSchedulerNode],
+]:
+    """
+    Compute buffer allocation and deallocation sizes and map their
+    lifetime to the node schedule
+    """
+
+    # get the execution step of each node, this will be used to determine
+    # the end_step of buffers
+    node_to_step: dict[BaseSchedulerNode, int] = {
+        node: step for step, node in enumerate(nodes)
+    }
+
+    # get buffers' size and liveliness information
+    buf_info_list: list[BufferInfo] = []
+    buf_to_snode_last_use: dict[
+        FreeableInputBuffer | SchedulerBuffer, BaseSchedulerNode
+    ] = {}
+
+    def _get_end_step_and_snode(
+        buf: FreeableInputBuffer | SchedulerBuffer,
+    ) -> tuple[int, BaseSchedulerNode | None]:
+        max_step: int = -1
+        max_step_snode: BaseSchedulerNode | None = None
+        succ_nodes = buf.mpi_buffer.succ_nodes
+        if succ_nodes:
+            for succ_node in succ_nodes:
+                step = node_to_step[succ_node]
+                if step > max_step:
+                    max_step = step
+                    max_step_snode = succ_node
+            assert max_step_snode is not None
+        return max_step, max_step_snode
+
+    # 1. for freeable input buffers
+    for buf_name, input_buf in name_to_freeable_input_buf.items():
+        end_step = -1
+        if buf_name not in graph_outputs:
+            end_step, end_step_snode = _get_end_step_and_snode(input_buf)
+            assert end_step_snode is not None
+            buf_to_snode_last_use[input_buf] = end_step_snode
+
+        buf_info_list.append(
+            BufferInfo(
+                input_buf,
+                input_buf.mpi_buffer.size_free,
+                input_buf.mpi_buffer.size_free,
+                0,
+                end_step,
+            )
+        )
+
+    # 2. for scheduler buffers
+    for step, node in enumerate(nodes):
+        for sched_buf in node.get_outputs():
+            # note: it is possible for a non-graph-output sched_buf to have no succ_nodes and
+            # to be only used by its defining op (e.g., due to fusion when all consumers of
+            # the buffer are fused with its defining op). In such cases, end_step is step.
+            buf_name = sched_buf.get_name()
+            end_step = -1
+            if buf_name not in graph_outputs:
+                end_step, end_step_snode = _get_end_step_and_snode(sched_buf)
+                if end_step == -1:
+                    end_step = step
+                    buf_to_snode_last_use[sched_buf] = node
+                else:
+                    assert end_step_snode is not None
+                    buf_to_snode_last_use[sched_buf] = end_step_snode
+
+            buf_info_list.append(
+                BufferInfo(
+                    sched_buf,
+                    sched_buf.mpi_buffer.size_alloc,
+                    sched_buf.mpi_buffer.size_free,
+                    step,
+                    end_step,
+                )
+            )
+
+    return buf_info_list, node_to_step, buf_to_snode_last_use
+
+
+def estimate_peak_memory(
+    nodes: list[BaseSchedulerNode],
+    name_to_freeable_input_buf: dict[str, FreeableInputBuffer],
+    graph_outputs: OrderedSet[str],
+) -> tuple[int, list[int]]:
+    """
+    Given a list of nodes in their execution order, estimate the peak memory, by
+    keeping track of the liveliness of SchedulerBuffers and FreeableInputBuffers.
+
+    Returns:
+        int: peak memory
+        List[int]: memory usage at each node (or each step).
+    """
+
+    buf_info_list, _, _ = compute_memory_timeline(
+        nodes, name_to_freeable_input_buf, graph_outputs
+    )
+
+    # incremental memory changes at each step
+    memory = [0 for _ in range(len(nodes) + 1)]
+
+    # for each buffer, update memory when created and when freed
+    for buf_info in buf_info_list:
+        memory[buf_info.start_step] += buf_info.size_alloc
+        memory[buf_info.end_step + 1] -= buf_info.size_free
+
+    # get peak memory by compute the cumulative memories
+    max_memory = 0
+    cur_memory = 0
+    memories_at_nodes = []
+    for t in range(len(nodes) + 1):
+        cur_memory += memory[t]
+        memories_at_nodes.append(cur_memory)
+        max_memory = max(max_memory, cur_memory)
+
+    return (max_memory, memories_at_nodes)
+
+
+@dataclasses.dataclass
+class SNodeMemory:
+    size_alloc: int
+    size_free: int
+
+
+def estimate_peak_memory_allocfree(
+    nodes: list[BaseSchedulerNode],
+    name_to_freeable_input_buf: dict[str, FreeableInputBuffer],
+    graph_outputs: OrderedSet[str],
+) -> tuple[
+    int,
+    list[tuple[int, int]],
+    dict[BaseSchedulerNode, SNodeMemory],
+    dict[FreeableInputBuffer | SchedulerBuffer, BaseSchedulerNode],
+]:
+    """
+    Alternative version of estimate_peak_memory, that respects the fact,
+    that every SchedulerNode has multiple phases:
+    1. alloc ( outputs )
+    2. run_kernel
+    3. dealloc last_use buffers
+    estimate_peak_memory collapses memory into one value: size_alloc - size_free
+    While peak memory happens after alloc.
+
+    Duplicating the code to not migrate all callsites at once,
+    In future usages of estimate_peak_memory will migrate to this version.
+    """
+
+    buf_info_list, _, buf_to_snode_last_use = compute_memory_timeline(
+        nodes, name_to_freeable_input_buf, graph_outputs
+    )
+
+    # incremental memory changes at each step
+    step_idx_allocfree = [SNodeMemory(0, 0) for _ in range(len(nodes))]
+
+    # for each buffer, update memory when created and when freed
+    for buf_info in buf_info_list:
+        step_idx_allocfree[buf_info.start_step].size_alloc += buf_info.size_alloc
+        if buf_info.end_step != -1:
+            step_idx_allocfree[buf_info.end_step].size_free += buf_info.size_free
+
+    snodes_allocfree = {}
+    for i, node in enumerate(nodes):
+        snodes_allocfree[node] = step_idx_allocfree[i]
+
+    max_memory = 0
+    cur_memory = 0
+    snodes_curr_memory = []
+    for t in range(len(nodes)):
+        alloc = step_idx_allocfree[t].size_alloc
+        free = step_idx_allocfree[t].size_free
+        cur_memory += alloc
+        post_alloc = cur_memory
+        max_memory = max(max_memory, cur_memory)
+        cur_memory -= free
+        post_free = cur_memory
+        snodes_curr_memory.append((post_alloc, post_free))
+
+    return (
+        max_memory,
+        snodes_curr_memory,
+        snodes_allocfree,
+        buf_to_snode_last_use,
+    )
+
+
+def topological_sort_lpmf(
+    nodes: list[BaseSchedulerNode],
+    name_to_freeable_input_buf: dict[str, FreeableInputBuffer],
+    name_to_buf: dict[str, SchedulerBuffer],
+    graph_outputs: OrderedSet[str],
+) -> list[BaseSchedulerNode]:
+    """
+    A bfs-based greedy topological order. LPMF stands for "Least Peak Memory First".
+
+    The idea is from this paper:
+    Buffer memory optimization for video codec application modeled in Simulink
+    https://www.cs.york.ac.uk/rts/docs/DAC-1964-2006/PAPERS/2006/DAC06/PDFFILES/P0689.PDF
+
+    The algorithm maintains the max memory so far.
+    At every iteration, for each scheduleable node, it computes:
+        - how much memory needs to be allocated for the output buffers of this node;
+        - how much memory can be freed as a result of executing this node.
+    This gives us two values for each node:
+        (1) mem1: memory during the execution of the node;
+        (2) mem2: memory after executing the node, after some input buffers are freed.
+    The greedy approach select as follows:
+        (i) if there are nodes whose mem1 values are below the max memory so far,
+            then pick the node with the lowest mem2 value;
+        (ii) otherwise, pick the one with the lowest mem1 value.
+    """
+
+    class NodeInfo(TypedDict):
+        indegree: int
+        memory_to_free: int
+
+    class BufferInfo(TypedDict):
+        outdegree: int
+
+    node_info: dict[BaseSchedulerNode, NodeInfo] = dict()
+    buf_info: dict[SchedulerBuffer | FreeableInputBuffer, BufferInfo] = dict()
+
+    # compute nodes' number of unmet dependencies (for schedulability)
+    # initialize the list of nodes ready to be scheduled
+    nodes_to_schedule: OrderedSet[BaseSchedulerNode] = OrderedSet()
+    for node in nodes:
+        node_info[node] = {
+            "indegree": len(node.mpi_node.pred_nodes),
+            "memory_to_free": 0,
+        }
+        if node_info[node]["indegree"] == 0:
+            nodes_to_schedule.add(node)
+
+    # compute buffers' number of unmet successors (used to decide when to free)
+    for buf in list(name_to_buf.values()) + list(name_to_freeable_input_buf.values()):
+        buf_info[buf] = {
+            "outdegree": len(buf.mpi_buffer.succ_nodes)
+            + (1 if buf.get_name() in graph_outputs else 0)
+        }
+
+    # initialize memory estimations
+    live_memory = sum(
+        input_buf.mpi_buffer.size_free
+        for input_buf in name_to_freeable_input_buf.values()
+    )
+
+    # this is the total output memory, which is a lower bound for peak memory
+    # we do not include the memory of non freeable input buffers
+    output_memory = 0
+    for buf_name in graph_outputs:
+        if buf_name in name_to_buf:
+            output_memory += name_to_buf[buf_name].mpi_buffer.size_free
+        elif buf_name in name_to_freeable_input_buf:
+            output_memory += name_to_freeable_input_buf[buf_name].mpi_buffer.size_free
+    max_memory = max(live_memory, output_memory)
+    memory_gap = max_memory - live_memory
+
+    # compute the amount of memory that is allocated when a node is scheduled
+    # and the amount of memory that can be freed when a node is scheduled
+    for node in nodes:
+        # 1. if a buffer read by this node is last used by this node
+        for buf in node.mpi_node.pred_buffers:
+            if buf_info[buf]["outdegree"] == 1:
+                node_info[node]["memory_to_free"] += buf.mpi_buffer.size_free
+        # 2. if a buffer written by this node is used internally and not used later
+        for buf in node.get_outputs():
+            if buf_info[buf]["outdegree"] == 0:
+                node_info[node]["memory_to_free"] += buf.mpi_buffer.size_free
+
+    # schedule nodes one at a time
+    schedule: list[BaseSchedulerNode] = []
+    size_threshold = config.size_threshold_for_succ_based_strategy
+    num_iters: int = 0
+    while num_iters < len(nodes) and nodes_to_schedule:
+        # select a node to schedule:
+        if (
+            size_threshold > 0
+            and min(node.mpi_node.size for node in nodes_to_schedule) > size_threshold
+        ):
+            selected_node = min(
+                nodes_to_schedule,
+                key=lambda node: min(
+                    (
+                        succ_node.mpi_node.index
+                        for succ_node in node.mpi_node.succ_nodes
+                    ),
+                    default=len(nodes),
+                ),
+            )
+        else:
+            selected_node = min(
+                nodes_to_schedule,
+                key=lambda node: (
+                    node.mpi_node.size if node.mpi_node.size > memory_gap else 0,
+                    node.mpi_node.size - node_info[node]["memory_to_free"],
+                    node.mpi_node.index,
+                ),
+            )
+        nodes_to_schedule.remove(selected_node)
+        schedule.append(selected_node)
+        num_iters += 1
+
+        # update memory usage
+        live_memory += selected_node.mpi_node.size
+        max_memory = max(max_memory, live_memory)
+        live_memory -= node_info[selected_node]["memory_to_free"]
+        memory_gap = max_memory - live_memory
+
+        # update successor nodes and nodes_to_schedule
+        for succ_node in selected_node.mpi_node.succ_nodes:
+            assert node_info[succ_node]["indegree"] > 0
+            node_info[succ_node]["indegree"] -= 1
+            if node_info[succ_node]["indegree"] == 0:
+                nodes_to_schedule.add(succ_node)
+
+        # update predecessor nodes
+        for buf in selected_node.mpi_node.pred_buffers:
+            assert buf_info[buf]["outdegree"] > 0
+            buf_info[buf]["outdegree"] -= 1
+            if buf_info[buf]["outdegree"] == 1:
+                for succ_node in buf.mpi_buffer.succ_nodes:
+                    node_info[succ_node]["memory_to_free"] += buf.mpi_buffer.size_free
+
+    if num_iters > len(nodes):
+        raise RuntimeError("Failed to schedule, while loop ran too long for lpmf")
+
+    return schedule
+
+
+def topological_sort_bfs(nodes: list[BaseSchedulerNode]) -> list[BaseSchedulerNode]:
+    """
+    A BFS topological sort that selects nodes whose dependencies are executed the
+    earliest. This follows a FIFO idea. Specifically, at every iteration, for each node
+    that is schedulable, we gather the order in which its predecessor nodes are executed,
+    and this sorted list of execution orders of predecessor nodes defines the priority.
+    We select the node whose predecessors nodes are executed the earliest. The FIFO
+    idea aims to reduce the liveness duration of buffers created.
+    """
+
+    class NodeInfo(TypedDict):
+        indegree: int
+        order: int
+
+    node_info: dict[BaseSchedulerNode, NodeInfo] = dict()
+
+    @dataclasses.dataclass
+    class NodeWithPriority:
+        priority: list[int]
+        node: BaseSchedulerNode
+
+        def __lt__(self, other: NodeWithPriority) -> bool:
+            if self.priority == other.priority:
+                return self.node.mpi_node.index < other.node.mpi_node.index
+            return self.priority < other.priority
+
+    def _node_priority(node: BaseSchedulerNode) -> list[int]:
+        # priority is the order in which predecessor nodes are executed
+        assert node_info[node]["indegree"] == 0
+        exec_orders = sorted(
+            OrderedSet(
+                node_info[pred_node]["order"] for pred_node in node.mpi_node.pred_nodes
+            )
+        )
+        return exec_orders
+
+    # compute nodes' number of unmet dependencies (for schedulability)
+    # initialize the list of nodes ready to be scheduled
+    nodes_to_schedule: list[NodeWithPriority] = []
+    for node in nodes:
+        node_info[node] = {"indegree": len(node.mpi_node.pred_nodes), "order": -1}
+        if node_info[node]["indegree"] == 0:
+            heapq.heappush(
+                nodes_to_schedule, NodeWithPriority(_node_priority(node), node)
+            )
+
+    # schedule nodes one at a time
+    schedule: list[BaseSchedulerNode] = []
+    num_iters: int = 0
+    while num_iters < len(nodes) and nodes_to_schedule:
+        # select a node to schedule
+        selected_node = heapq.heappop(nodes_to_schedule).node
+        node_info[selected_node]["order"] = len(schedule)
+        schedule.append(selected_node)
+        num_iters += 1
+
+        # update successor nodes and nodes_to_schedule
+        for succ_node in selected_node.mpi_node.succ_nodes:
+            assert node_info[succ_node]["indegree"] > 0
+            node_info[succ_node]["indegree"] -= 1
+            if node_info[succ_node]["indegree"] == 0:
+                heapq.heappush(
+                    nodes_to_schedule,
+                    NodeWithPriority(_node_priority(succ_node), succ_node),
+                )
+
+    if num_iters > len(nodes):
+        raise RuntimeError("Failed to schedule, while loop ran too long for bfs")
+
+    return schedule
+
+
+def topological_sort_dfs(nodes: list[BaseSchedulerNode]) -> list[BaseSchedulerNode]:
+    """
+    This is a DFS topological sort. The setup is similar to `topological_sort_schedule`
+    in scheduler.py. The difference is the order nodes are visited in the outer loop.
+    In `topological_sort_schedule`, nodes are visited in their original order.
+    In this function, nodes are visited based on their priority -- for each node, we
+    compute the total memory of all buffers it reads from or writes to, and we visit
+    the nodes in ascending order of this priority.
+    """
+    seen: OrderedSet[BaseSchedulerNode] = OrderedSet()
+    name_to_node: dict[str, BaseSchedulerNode] = dict()
+    result: list[BaseSchedulerNode] = []
+    size_with_reads: dict[BaseSchedulerNode, int] = dict()
+
+    def visit(n: BaseSchedulerNode) -> None:
+        if n not in seen:
+            seen.add(n)
+            dep_nodes = [
+                name_to_node[dep.name]
+                for dep in n.unmet_dependencies
+                if dep.name in name_to_node
+            ]
+            for node in sorted(
+                dep_nodes, key=lambda n: (size_with_reads[n], n.mpi_node.index)
+            ):
+                visit(node)
+            result.append(n)
+
+    for node in nodes:
+        for name in node.get_buffer_names():
+            name_to_node[name] = node
+
+    for node in nodes:
+        size_with_reads[node] = node.mpi_node.size + sum(
+            pred_buf.mpi_buffer.size_free for pred_buf in node.mpi_node.pred_buffers
+        )
+    for node in sorted(nodes, key=lambda n: (size_with_reads[n], n.mpi_node.index)):
+        visit(node)
+
+    return result
+
+
+def validate_graph_acyclic(nodes: list[BaseSchedulerNode]) -> None:
+    """
+    Validate that the graph is acyclic by checking predecessor relationships.
+
+    Raises:
+        RuntimeError: If a cycle is detected in the graph
+    """
+    # DFS coloring scheme for cycle detection:
+    # WHITE (0): Node has not been visited yet
+    # GRAY (1): Node is currently being processed (in the recursion stack)
+    # BLACK (2): Node has been completely processed (finished exploring all its predecessors)
+    # A back edge (cycle) is detected when we encounter a GRAY node during DFS traversal
+    WHITE, GRAY, BLACK = 0, 1, 2
+    color = dict.fromkeys(nodes, WHITE)
+    path: list[BaseSchedulerNode] = []  # Track current DFS path
+
+    def dfs_visit(node: BaseSchedulerNode) -> None:
+        if color[node] == BLACK:
+            return
+
+        if color[node] == GRAY:
+            path.append(node)
+            path_info = " -> ".join([node.get_name() for node in path])
+
+            raise RuntimeError(
+                f"Cycle detected in memory planning graph"
+                f"Path containing cycle (i -> j: j is a dependency of i): {path_info} "
+                f"This indicates invalid dependency relationships in the scheduler graph"
+            )
+
+        color[node] = GRAY
+        path.append(node)
+
+        for pred_node in node.mpi_node.pred_nodes:
+            assert pred_node != node
+            dfs_visit(pred_node)
+
+        path.pop()
+        color[node] = BLACK
+
+    # Start DFS from all unvisited nodes
+    for node in nodes:
+        if color[node] == WHITE:
+            dfs_visit(node)
+
+
+def validate_unique_buffer_names(
+    nodes: list[BaseSchedulerNode],
+    name_to_buf: dict[str, SchedulerBuffer],
+    name_to_freeable_input_buf: dict[str, FreeableInputBuffer],
+) -> None:
+    """
+    Validate that for each node's output buffer, the name_to_buf mapping is correct.
+    For each output buffer buf, we should have name_to_buf[buf.get_name()] == buf.
+    Also validate that no buffer names overlap with freeable input buffer names.
+
+    Raises:
+        RuntimeError: If buffer name mapping is incorrect or names overlap
+    """
+    for node in nodes:
+        for buf in node.get_outputs():
+            buf_name = buf.get_name()
+
+            # Check if buffer name exists in the mapping
+            if buf_name not in name_to_buf:
+                raise RuntimeError(
+                    f"{buf_name} from {node.get_name()} is not found in name_to_buf mapping."
+                    f" This indicates a missing buffer mapping."
+                )
+
+            # Check if the mapping points to the correct buffer object
+            if name_to_buf[buf_name] != buf:
+                raise RuntimeError(
+                    f"Buffer name mapping is incorrect for '{buf_name}'."
+                    f"Expected name_to_buf['{buf_name}'] to be {buf.debug_str()}"
+                    f"but got {name_to_buf[buf_name].debug_str()}"
+                    f"This indicates some buffers share the same name"
+                )
+
+            # Check if buffer name conflicts with freeable input buffer names
+            if buf_name in name_to_freeable_input_buf:
+                raise RuntimeError(
+                    f"Buffer name conflict detected: '{buf_name}' from node {node.get_name()} "
+                    f"is also used as a freeable input buffer name. "
+                )
+
+
+def prepare_planning_info(
+    nodes: list[BaseSchedulerNode],
+    name_to_buf: dict[str, SchedulerBuffer],
+    name_to_fused_node: dict[str, BaseSchedulerNode],
+    graph_inputs: OrderedSet[str],
+    graph_outputs: OrderedSet[str],
+) -> tuple[int, dict[str, FreeableInputBuffer]]:
+    """
+    Prepare planning info. As nodes are scheduled one at a time, these help
+    keep track of when a buffer can be freed, and when a node can be scheduled
+
+    Returns:
+        int: peak memory estimation
+        dict[str, FreeableInputBuffer]: name to freeable input buffer
+    """
+    name_to_freeable_input_buf = get_freeable_input_buf(nodes, graph_inputs)
+    assign_memory_planning_info_for_scheduler_buffers(nodes, name_to_buf)
+    assign_memory_planning_info_for_scheduler_nodes(
+        nodes, name_to_fused_node, name_to_buf, name_to_freeable_input_buf
+    )
+
+    # the default
+    estimated_peak_memory, _ = estimate_peak_memory(
+        nodes, name_to_freeable_input_buf, graph_outputs
+    )
+
+    return estimated_peak_memory, name_to_freeable_input_buf
+
+
+def reorder_for_peak_memory(
+    nodes: list[BaseSchedulerNode],
+    name_to_buf: dict[str, SchedulerBuffer],
+    name_to_fused_node: dict[str, BaseSchedulerNode],
+    graph_inputs: OrderedSet[str],
+    graph_outputs: OrderedSet[str],
+    methods: list[Callable[..., list[BaseSchedulerNode]]] = [  # noqa: B006
+        topological_sort_lpmf,
+        topological_sort_bfs,
+        topological_sort_dfs,
+    ],
+) -> list[BaseSchedulerNode]:
+    """
+    Try a few heuristics based topological sort algorithms, and pick the one whose
+    resulting topological order has the lowest peak memory estimation.
+    """
+
+    torch_log.info("Reordering for peak memory -- %d nodes", len(nodes))
+
+    estimated_peak_memory, name_to_freeable_input_buf = prepare_planning_info(
+        nodes,
+        name_to_buf,
+        name_to_fused_node,
+        graph_inputs,
+        graph_outputs,
+    )
+
+    # export graph for simulator if needed
+    if config.reorder_for_peak_memory_debug:
+        export_graph_for_simulator(
+            nodes,
+            name_to_freeable_input_buf,
+            name_to_fused_node,
+            graph_inputs,
+            graph_outputs,
+        )
+
+    # Validate planning info before proceeding with reordering
+    try:
+        validate_graph_acyclic(nodes)
+        validate_unique_buffer_names(nodes, name_to_buf, name_to_freeable_input_buf)
+    except RuntimeError:
+        torch_log.exception("Memory planning validation failed")
+        if not is_fbcode():  # TODO: remove after ensuring OSS side is safe
+            raise
+
+    # keep track of the peak memory estimates of different methods
+    peak_memory_diff_methods: list[PeakMemoryResult] = []
+    peak_memory_diff_methods.append(
+        PeakMemoryResult(nodes, estimated_peak_memory, "baseline")
+    )
+    torch_log.info("Baseline peak memory: %d", estimated_peak_memory)
+
+    # other methods
+    for method in methods:
+        try:
+            if method is topological_sort_lpmf:
+                order = method(
+                    nodes, name_to_freeable_input_buf, name_to_buf, graph_outputs
+                )
+            else:
+                order = method(nodes)
+            assert len(order) == len(nodes)
+            peak_memory, _ = estimate_peak_memory(
+                order, name_to_freeable_input_buf, graph_outputs
+            )
+            peak_memory_diff_methods.append(
+                PeakMemoryResult(order, peak_memory, method.__name__)
+            )
+            torch_log.info("%s peak memory: %d", method.__name__, peak_memory)
+        except Exception:
+            torch_log.exception("Failed to reorder for %s", method.__name__)
+            if not is_fbcode():  # TODO: remove after ensuring OSS side is safe
+                raise
+
+    signpost_event(
+        category="inductor",
+        name="memory",
+        parameters={
+            "orm": {elem.method: elem.peak_memory for elem in peak_memory_diff_methods},
+        },
+    )
+
+    # get the optimal one
+    best_result = min(peak_memory_diff_methods, key=lambda x: x.peak_memory)
+
+    return best_result.order
+
+
+def export_graph_for_simulator(
+    nodes: list[BaseSchedulerNode],
+    name_to_freeable_input_buf: dict[str, FreeableInputBuffer],
+    name_to_fused_node: dict[str, BaseSchedulerNode],
+    graph_inputs: OrderedSet[str],
+    graph_outputs: OrderedSet[str],
+) -> None:
+    """
+    This is for debugging purposes. It will dump a json file that records graph information.
+    The graph can then be used in a simulator: https://fburl.com/code/3l3d3qi4
+    """
+
+    class ORMBuffer(TypedDict):
+        name: str
+        size_alloc: int
+        size_free: int
+        size: int  # for backward compatibility
+        is_input: bool
+        is_output: bool
+        deps: list[str]
+        unmet_deps: list[str]
+
+    class ORMNode(TypedDict):
+        name: str
+        buffer_names: list[str]
+
+    class ORMGraph(TypedDict):
+        nodes: list[ORMNode]
+        buffers: list[ORMBuffer]
+
+    orm_buffers: list[ORMBuffer] = []
+    orm_nodes: list[ORMNode] = []
+
+    # get orm buffers for freeable input buffers
+    for buf_name, input_buf in name_to_freeable_input_buf.items():
+        orm_buf_input_buffer: ORMBuffer = {
+            "name": buf_name,
+            "size_alloc": input_buf.mpi_buffer.size_free,
+            "size_free": input_buf.mpi_buffer.size_free,
+            "size": input_buf.mpi_buffer.size_free,
+            "is_input": True,
+            "is_output": buf_name in graph_outputs,
+            "deps": [],
+            "unmet_deps": [],
+        }
+        orm_buffers.append(orm_buf_input_buffer)
+
+    # get orm buffers for scheduler buffers
+    name_to_buf: dict[str, SchedulerBuffer] = {
+        buf.get_name(): buf for node in nodes for buf in node.get_outputs()
+    }  # need to reassign due to probably node pruning
+    for buf_name, sched_buf in name_to_buf.items():
+        if sched_buf.defining_op is None:
+            continue
+        deps = [
+            pred_buf.get_name()
+            for pred_buf in name_to_fused_node[
+                sched_buf.defining_op.get_name()
+            ].mpi_node.pred_buffers
+        ]
+        orm_buf_scheduler_buffer: ORMBuffer = {
+            "name": buf_name,
+            "size_alloc": sched_buf.mpi_buffer.size_alloc,
+            "size_free": sched_buf.mpi_buffer.size_free,
+            "size": sched_buf.mpi_buffer.size_free,
+            "is_input": False,
+            "is_output": buf_name in graph_outputs,
+            "deps": deps,
+            "unmet_deps": [
+                buf_name for buf_name in deps if buf_name not in graph_inputs
+            ],
+        }
+        orm_buffers.append(orm_buf_scheduler_buffer)
+
+    # get orm nodes
+    for node in nodes:
+        orm_node: ORMNode = {
+            "name": node.get_name(),
+            "buffer_names": list(node.get_buffer_names()),
+        }
+        orm_nodes.append(orm_node)
+
+    # create the graph object
+    g: ORMGraph = {
+        "nodes": orm_nodes,
+        "buffers": orm_buffers,
+    }
+
+    # dump the graph
+    import json
+    import os
+
+    import torch
+    from functorch.compile import get_graph_being_compiled
+
+    name = os.path.splitext(get_graph_being_compiled())[0] + "_fused"
+
+    g_str = json.dumps(g, indent=2)
+
+    torch._logging.trace_structured(
+        "artifact",
+        metadata_fn=lambda: {
+            "name": name,
+            "encoding": "string",
+        },
+        payload_fn=lambda: g_str,
+    )

outputs/audit_venv/lib/python3.11/site-packages/torch/_inductor/metrics.py

@@ -0,0 +1,484 @@
+from __future__ import annotations
+
+import csv
+import dataclasses
+import inspect
+import os
+import re
+from dataclasses import dataclass
+from functools import lru_cache
+from typing import TYPE_CHECKING
+
+from torch._inductor import config
+from torch._inductor.utils import get_benchmark_name
+from torch.utils._ordered_set import OrderedSet
+
+
+# Prevent circular import
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from torch._inductor.runtime.triton_compat import Config
+    from torch._inductor.scheduler import BaseSchedulerNode
+
+# counter for tracking how many kernels have been generated
+generated_kernel_count = 0
+generated_cpp_vec_kernel_count = 0
+num_bytes_accessed = 0
+nodes_num_elem: list[
+    tuple[
+        BaseSchedulerNode,
+        int,
+    ]
+] = []
+node_runtimes: list[tuple[BaseSchedulerNode, float]] = []
+
+# counters for tracking fusions
+ir_nodes_pre_fusion = 0
+
+# counters for tracking to_dtype inserted
+cpp_to_dtype_count = 0
+
+
+@dataclasses.dataclass
+class CppOuterLoopFusedCount:
+    inner_kernel_number: int
+    local_buffer_number: int = 0
+
+
+# The length counts the number of outer loop fusions.
+cpp_outer_loop_fused_inner_counts: list[CppOuterLoopFusedCount] = []
+
+num_comprehensive_padding = 0
+num_matches_for_scatter_upon_const_tensor = 0
+
+num_loop_reordering = 0
+num_auto_chunking: int = 0
+
+# counter for parallel reduction.
+parallel_reduction_count = 0
+
+codegen_mix_order_reduction = 0
+
+
+# reset all counters
+def reset() -> None:
+    global generated_kernel_count
+    global generated_cpp_vec_kernel_count
+    global num_bytes_accessed, nodes_num_elem
+    global ir_nodes_pre_fusion
+    global cpp_to_dtype_count
+    global cpp_outer_loop_fused_inner_counts
+    global num_comprehensive_padding
+    global num_matches_for_scatter_upon_const_tensor
+    global num_loop_reordering
+    global parallel_reduction_count
+    global codegen_mix_order_reduction
+    global num_auto_chunking
+
+    generated_kernel_count = 0
+    generated_cpp_vec_kernel_count = 0
+    num_bytes_accessed = 0
+    nodes_num_elem.clear()
+    node_runtimes.clear()
+    ir_nodes_pre_fusion = 0
+    cpp_to_dtype_count = 0
+    cpp_outer_loop_fused_inner_counts.clear()
+    num_comprehensive_padding = 0
+    num_matches_for_scatter_upon_const_tensor = 0
+    num_loop_reordering = 0
+    parallel_reduction_count = 0
+    codegen_mix_order_reduction = 0
+    num_auto_chunking = 0
+
+
+@dataclass
+class CachedMetricsDeltas:
+    """
+    The subset of metrics we want update across cache hits, e.g., the
+    FxGraphCache.
+    """
+
+    generated_kernel_count: int
+    generated_cpp_vec_kernel_count: int
+    ir_nodes_pre_fusion: int
+    cpp_to_dtype_count: int
+    num_bytes_accessed: int
+    num_matches_for_scatter_upon_const_tensor: int
+
+
+def get_metric_fields() -> list[str]:
+    return [field.name for field in dataclasses.fields(CachedMetricsDeltas)]
+
+
+class CachedMetricsHelper:
+    """
+    A helper class to help calculate and apply counter deltas for those
+    metrics we want to save with cache entries (e.g., FxGraphCache) and
+    apply on a cache hit.
+    """
+
+    def __init__(self) -> None:
+        self.cached_metrics = {}
+        for metric in get_metric_fields():
+            self.cached_metrics[metric] = globals()[metric]
+
+    def get_deltas(self) -> CachedMetricsDeltas:
+        delta_metrics = {}
+        for metric in get_metric_fields():
+            delta_metrics[metric] = globals()[metric] - self.cached_metrics[metric]
+
+        return CachedMetricsDeltas(**delta_metrics)
+
+    @staticmethod
+    def apply_deltas(delta: CachedMetricsDeltas) -> None:
+        for metric in get_metric_fields():
+            globals()[metric] += getattr(delta, metric)
+
+
+REGISTERED_METRIC_TABLES: dict[str, MetricTable] = {}
+
+
+@dataclass
+class MetricTable:
+    table_name: str
+    column_names: list[str]
+
+    num_rows_added: int = 0
+
+    def add_row(self, row_fn: Callable[[], dict[str, str | float | None]]) -> None:
+        if self.table_name not in enabled_metric_tables():
+            return
+
+        row_dict = row_fn()
+        assert len(self.column_names) == len(row_dict), (
+            f"{len(self.column_names)} v.s. {len(row_dict)}"
+        )
+        assert OrderedSet(self.column_names) == OrderedSet(row_dict.keys()), (
+            f"{OrderedSet(self.column_names)} v.s. {OrderedSet(row_dict.keys())}"
+        )
+
+        bn = get_benchmark_name()
+        # assert bn is not None
+        row = [bn] + [row_dict[column_name] for column_name in self.column_names]
+        assert all(isinstance(i, (str, float, type(None))) for i in row)
+        self._write_row(row)
+
+    def output_filename(self) -> str:
+        return f"metric_table_{self.table_name}.csv"
+
+    def write_header(self) -> None:
+        filename = self.output_filename()
+        with open(filename, "w") as fd:
+            writer = csv.writer(fd, lineterminator="\n")
+            writer.writerow(["model_name"] + self.column_names)
+
+    def _write_row(self, row: list[str | float | None]) -> None:
+        filename = self.output_filename()
+        if self.num_rows_added == 0 and not os.path.exists(filename):
+            self.write_header()
+
+        self.num_rows_added += 1
+
+        for idx, orig_val in enumerate(row):
+            if isinstance(orig_val, float):
+                new_val = f"{orig_val:.6f}"
+            elif orig_val is None:
+                new_val = ""
+            else:
+                new_val = orig_val
+            row[idx] = new_val
+
+        with open(filename, "a") as fd:
+            writer = csv.writer(fd, lineterminator="\n")
+            writer.writerow(row)
+
+    @staticmethod
+    def register_table(name: str, column_names: list[str]) -> None:
+        table = MetricTable(name, column_names)
+        REGISTERED_METRIC_TABLES[name] = table
+
+
+MetricTable.register_table(
+    "slow_fusion",
+    [
+        "kernel1_path",
+        "kernel1_latency",
+        "kernel2_path",
+        "kernel2_latency",
+        "fused_kernel_path",
+        "fused_kernel_latency",
+        "slow_down_ratio",
+    ],
+)
+
+# track the fusion statistics for each graph
+MetricTable.register_table(
+    "graph_stats",
+    [
+        "graph_id",
+        "num_nodes_before_fusion",
+        "num_nodes_after_fusion",
+    ],
+)
+
+# track the perf difference between persistent reduction and non-persistent
+# reductions
+MetricTable.register_table(
+    "persistent_red_perf",
+    [
+        "kernel0_path",
+        "kernel1_path",
+        "kernel2_path",
+        "kernel3_path",
+        "kernel0_latency",
+        "kernel1_latency",
+        "kernel2_latency",
+        "kernel3_latency",
+        "size_hints",
+        "reduction_hint",
+    ],
+)
+
+# Log the fusion failures due to indexing mismatch
+MetricTable.register_table(
+    "fusion_failure_due_to_indexing_mismatch",
+    [
+        "pre_grad_graph_id",
+        "post_grad_graph_id",
+        "node1_name",
+        "node2_name",
+        "node1_debug_str",
+        "node2_debug_str",
+        "common_buffer_names",
+        "failure_reason",
+    ],
+)
+
+# Log metadata for pointwise/reduction kernels. E.g., model name, kernel path, numel, rnumel, reduction hint
+MetricTable.register_table(
+    "kernel_metadata",
+    [
+        "kernel_name",
+        "kernel_path",
+        "kernel_category",  # pointwise/reduction/foreach etc.
+        "size_hints",
+        "reduction_hint",
+        "line_of_code",
+        "num_load",
+        "num_store",
+        "num_for_loop",
+        "num_atomic_add",
+        "num_args",
+        # xyz numel can be different to size_hints since size_hints are rounded
+        # up to the nearest power of 2.
+        # Inductor kernel will burn in the xyz numel in kernel code for static
+        # shape kernels.
+        # Logging them will be helpful to find unaligned shape for reduction
+        "xnumel",
+        "ynumel",
+        "rnumel",
+        "kernel_args_num_gb",
+    ],
+)
+
+
+def _parse_kernel_fn_code(kernel_module_code: str) -> str:
+    """
+    The kernel_module_code is the python module that contains kernel function code.
+    kernel function is the proper triton kernel function annotated with
+    @triton.jit
+    """
+    from .codecache import PyCodeCache
+    from .wrapper_benchmark import get_triton_kernel
+
+    mod = PyCodeCache.load(kernel_module_code)
+    kernel = get_triton_kernel(mod)
+    # kernel is a CachingAutotune; kernel.fn is the JITFunction;
+    # kernel.fn.fn is the function being decorate by triton.jit
+    return inspect.getsource(kernel.fn.fn)
+
+
+def _parse_kernel_line_of_code(proper_kernel_fn_code: str) -> int:
+    """
+    Return the line of code for the kernel excluding the decorators.
+    """
+    return len(proper_kernel_fn_code.splitlines())
+
+
+def _parse_size_hints(kernel_module_code: str, kernel_category: str) -> str | None:
+    if kernel_category == "foreach":
+        # foreach kernel does not have size_hints
+        return None
+    m = re.search(r"size_hints=(\[[0-9, ]*\]),", kernel_module_code)
+    assert m, "size_hints missing!"
+    return m.group(1)
+
+
+def _parse_reduction_hint(kernel_category: str, kernel_module_code: str) -> str | None:
+    if kernel_category not in ("reduction", "persistent_reduction"):
+        return None
+    m = re.search(r"reduction_hint=ReductionHint\.(\w*),", kernel_module_code)
+    assert m, "reduction_hint not found in kernel source code!"
+    return m.group(1)
+
+
+def _count_pattern(proper_kernel_fn_code: str, pattern: str) -> int:
+    return proper_kernel_fn_code.count(pattern)
+
+
+def _count_args(proper_kernel_fn_code: str) -> int:
+    def_line = proper_kernel_fn_code.splitlines()[0]
+    assert def_line.startswith("def ")
+    start_idx = def_line.index("(")
+    end_idx = def_line.index("):")
+    decl_csv = def_line[start_idx + 1 : end_idx]
+    comps = decl_csv.split(",")
+    return len(comps)
+
+
+def _parse_proper_kernel_fn_code(kernel_fn_code: str) -> str:
+    """
+    Skip decorators.
+    """
+    start_pos = kernel_fn_code.index("def ")
+    return kernel_fn_code[start_pos:]
+
+
+def _parse_numel(proper_kernel_fn_code: str, numel_arg_name: str) -> int | None:
+    m = re.search(f"{numel_arg_name} = ([\\d]+)", proper_kernel_fn_code)
+    if m:
+        return int(m.group(1))
+    else:
+        return None
+
+
+def _parse_kernel_args_num_gb(
+    kernel_fn_code: str, kernel_category: str
+) -> float | None:
+    """
+    inductor meta looks like:
+        inductor_meta={... 'mutated_arg_names': [], 'no_x_dim': False, 'kernel_num_gb': 2.0},
+    """
+    m = re.search(r".kernel_num_gb.:\s*([0-9.]+)", kernel_fn_code)
+    if m:
+        return float(m.group(1))
+    else:
+        """
+        There are a few cases that kernel_num_gdb field can be missing:
+        1. the field will be missing if config.benchmark_kernel and
+           config.profile_bandwidth are false
+        2. even if config.benchmark_kernel or config.profile_bandwidth is true.
+           foreach kernel does not have kernel_num_gb field in the metadata
+        """
+        return None
+
+
+def log_kernel_metadata(
+    kernel_name: str, kernel_path: str, kernel_module_code: str
+) -> None:
+    """
+    An utility to log kernel metadata. We may parse metadata from kernel source code here.
+
+    It's fine to parse the generated kernel code here since the logging is
+    disabled by default. It would hurt compilation time.
+    """
+    from .wrapper_benchmark import get_kernel_category_by_source_code
+
+    kernel_category = get_kernel_category_by_source_code(kernel_module_code)
+    reduction_hint = _parse_reduction_hint(kernel_category, kernel_module_code)
+    size_hints = _parse_size_hints(kernel_module_code, kernel_category)
+    kernel_fn_code = _parse_kernel_fn_code(kernel_module_code)
+
+    proper_kernel_fn_code = _parse_proper_kernel_fn_code(kernel_fn_code)
+
+    # the line of code excluding the decortors
+    kernel_line_of_code = _parse_kernel_line_of_code(proper_kernel_fn_code)
+
+    get_metric_table("kernel_metadata").add_row(
+        lambda: {
+            "kernel_name": kernel_name,
+            "kernel_path": kernel_path,
+            "kernel_category": kernel_category,
+            "size_hints": size_hints,
+            "reduction_hint": reduction_hint,
+            "line_of_code": kernel_line_of_code,
+            "num_load": _count_pattern(proper_kernel_fn_code, "tl.load"),
+            "num_store": _count_pattern(proper_kernel_fn_code, "tl.store"),
+            "num_for_loop": _count_pattern(proper_kernel_fn_code, "for "),
+            "num_atomic_add": _count_pattern(proper_kernel_fn_code, "tl.atomic_add"),
+            "num_args": _count_args(proper_kernel_fn_code),
+            "xnumel": _parse_numel(proper_kernel_fn_code, "xnumel"),
+            "ynumel": _parse_numel(proper_kernel_fn_code, "ynumel"),
+            "rnumel": _parse_numel(proper_kernel_fn_code, "rnumel"),
+            "kernel_args_num_gb": _parse_kernel_args_num_gb(
+                kernel_fn_code, kernel_category
+            ),
+        }
+    )
+
+
+def purge_old_log_files() -> None:
+    """
+    Purge the old log file at the beginning when the benchmark script runs.
+    Should do it in the parent process rather than the child processes running
+    each individual model.
+    """
+    for name, table in REGISTERED_METRIC_TABLES.items():
+        if name in enabled_metric_tables():
+            filename = table.output_filename()
+            if os.path.exists(filename):
+                os.unlink(filename)
+
+            table.write_header()
+
+
+def enabled_metric_tables() -> OrderedSet[str]:
+    return enabled_metric_tables_impl(config.enabled_metric_tables)
+
+
+@lru_cache
+def enabled_metric_tables_impl(config_str: str) -> OrderedSet[str]:
+    enabled: OrderedSet[str] = OrderedSet()
+    for name in config_str.split(","):
+        name = name.strip()
+        if not name:
+            continue
+        assert name in REGISTERED_METRIC_TABLES, (
+            f"Metric table name {name} is not registered"
+        )
+        enabled.add(name)
+    return enabled
+
+
+def is_metric_table_enabled(name: str) -> bool:
+    return name in enabled_metric_tables()
+
+
+def get_metric_table(name: str) -> MetricTable:
+    assert name in REGISTERED_METRIC_TABLES, f"Metric table {name} is not defined"
+    return REGISTERED_METRIC_TABLES[name]
+
+
+MetricTable.register_table(
+    "kernel_autotune",
+    [
+        "kernel_path",
+        "kernel_name",
+        "triton_config",
+        "latency_ms",
+    ],
+)
+
+
+def log_kernel_autotune_result(
+    kernel_path: str, kernel_name: str, config: Config, latency: float
+) -> None:
+    get_metric_table("kernel_autotune").add_row(
+        lambda: {
+            "kernel_path": kernel_path,
+            "kernel_name": kernel_name,
+            "triton_config": str(config),
+            "latency_ms": latency,
+        }
+    )

outputs/audit_venv/lib/python3.11/site-packages/torch/_inductor/mkldnn_ir.py

@@ -0,0 +1,1361 @@
+# mypy: allow-untyped-defs
+from collections.abc import Sequence
+from typing import Any, Optional
+
+import sympy
+
+import torch
+from torch._prims_common import make_channels_last_strides_for, StrideType
+from torch.utils._ordered_set import OrderedSet
+
+from .ir import (
+    ExternKernelAlloc,
+    FixedLayout,
+    FlexibleLayout,
+    get_device_type,
+    ir_node_to_tensor,
+    IRNode,
+    is_contiguous_storage_and_layout,
+    Layout,
+    may_convert_to_optional,
+    MultiOutput,
+    MultiOutputLayout,
+    MutationOutput,
+    NoneLayout,
+    TensorBox,
+)
+from .utils import convert_shape_to_inductor, pad_listlike, SUPPORTED_MKLDNN_DEVICES
+from .virtualized import V
+
+
+def _prepare_convolution_fusion_create(
+    cls,
+    x: "TensorBox",
+    weight: "TensorBox",
+    bias: "TensorBox",
+    padding: Sequence[int],
+    stride: Sequence[int],
+    dilation: Sequence[int],
+    groups: int,
+    transposed: bool = False,
+    output_padding: Sequence[int] | None = None,
+    quantize_args: list["TensorBox"] | None = None,
+    other: Optional["TensorBox"] = None,
+):
+    """
+    This function is a helper function to prepare inputs, layout and constant args
+    for convolution post-op fusion's create function, including deciding the output
+    layout (channels first or channels last), realizing inputs and make them etc. The
+    function only supports the CPU/XPU device since conv post-op fusion kernel is only
+    supported on CPU/XPU right now.
+    """
+
+    # Port from aten/src/ATen/native/ConvUtils.h: _conv_input_size
+    def _conv_input_size(
+        output_size, weight_size, padding, output_padding, stride, dilation, groups
+    ):
+        assert len(output_size) == len(weight_size), "Expect input dim == weight dim"
+        dim = len(output_size)
+        assert dim > 2, "Expect input dim > 2"
+
+        BATCH_DIM = 0
+        WEIGHT_INPUT_CHANNELS_DIM = 1
+        input_size = []
+        input_size.append(output_size[BATCH_DIM])
+        input_size.append(weight_size[WEIGHT_INPUT_CHANNELS_DIM] * groups)
+        for d in range(2, dim):
+            kernel = (weight_size[d] - 1) * dilation[d - 2] + 1
+            input_size_d = (
+                (output_size[d] - 1) * stride[d - 2]
+                - (padding[d - 2] * 2)
+                + kernel
+                + output_padding[d - 2]
+            )
+            input_size.append(input_size_d)
+        return list(map(int, input_size))
+
+    # Port from aten/src/ATen/native/ConvUtils.h: _conv_output_size
+    def _conv_output_size(input_size, weight_size, padding, stride, dilation=None):
+        has_dilation = dilation is not None
+        dim = len(input_size)
+        output_size = []
+        output_size.append(input_size[0])
+        output_size.append(weight_size[0])
+        for d in range(2, dim):
+            # pyrefly: ignore [unsupported-operation]
+            dilation_ = dilation[d - 2] if has_dilation else 1
+            kernel = dilation_ * (weight_size[d] - 1) + 1
+            output_size_d = (input_size[d] + (2 * padding[d - 2]) - kernel) // stride[
+                d - 2
+            ] + 1
+            output_size.append(output_size_d)
+        return output_size
+
+    # The size of prepacked_weight is the prepacked weight size of deconv:
+    #   Groups > 1:  [g*o, i/g, ...]
+    #   Groups == 1: [o, i, ...]
+    # Returns original weight size in [i, o, ...]
+    def _original_deconv_weight_size(
+        prepacked_weight,
+        groups,
+    ):
+        prepacked_weight_size = prepacked_weight.size()
+        dim = len(prepacked_weight_size)
+        assert dim > 2, "Expect weight dim > 2"
+        if groups > 1:
+            weight_size = []
+            weight_size.append(prepacked_weight_size[1] * groups)
+            weight_size.append(prepacked_weight_size[0] / groups)
+            weight_size.extend(prepacked_weight_size[d] for d in range(2, dim))
+        else:
+            weight_size = prepacked_weight.transpose(0, 1).size()
+        return weight_size
+
+    x.realize()
+    weight.realize()
+    if bias is not None:
+        bias.realize()
+    with V.graph.fake_mode:
+        # TODO <Leslie> cleaned up the fake_tensor trace as Linear implementation
+        x_fake = ir_node_to_tensor(x)
+        weight_fake = ir_node_to_tensor(weight)
+        dims = len(x_fake.size()) - 2
+        assert 0 < len(padding) <= dims
+        assert 0 < len(dilation) <= dims
+        assert 0 < len(stride) <= dims
+        padding = pad_listlike(padding, dims)
+        dilation = pad_listlike(dilation, dims)
+        stride = pad_listlike(stride, dims)
+        if output_padding is None:
+            output_padding = pad_listlike([0], dims)
+        else:
+            assert 0 < len(output_padding) <= dims
+            output_padding = pad_listlike(output_padding, dims)
+        assert isinstance(groups, (int, sympy.core.numbers.Integer))
+        if transposed:
+            # When transposed, the size of the prepacked oneDNN weight is different
+            # from the PyTorch weight. We're not able to run aten conv with such
+            # size. We infer the output size from the input params here:
+            weight_size = _original_deconv_weight_size(weight_fake, groups)
+            input_size = x_fake.size()
+            output_size = _conv_input_size(
+                input_size,
+                weight_size,
+                padding,
+                output_padding,
+                stride,
+                dilation,
+                groups,
+            )
+        else:
+            x_shape = list(x_fake.shape)
+            weight_shape = list(weight_fake.shape)
+            if len(x_shape) != len(weight_shape):
+                assert len(x_shape) == 3 and len(weight_shape) == 4
+                weight_shape.pop(2)
+            output_size = _conv_output_size(
+                x_shape,
+                weight_shape,
+                padding,
+                stride,
+                dilation,
+            )
+
+        req_stride_order = [0] + list(reversed(range(1, len(stride) + 1)))
+        req_stride_order = [len(req_stride_order)] + req_stride_order
+
+    x = cls.require_stride_order(x, req_stride_order)
+
+    # We won't do weight prepack for Conv if dynamic_shapes or if is xpu.
+    # In static shape cases, since weight is prepacked, we'll always force output to be channels last in the Conv kernel.
+    # In dynamic shape cases, for input with channels = 1, like tensor of size (s0, 1, 28, 28) and stride (784, 784, 28, 1),
+    # x = cls.require_stride_order(x, req_stride_order) where req_stride_order is in the channels last order
+    # won't change the stride of this tensor since stride for dimensions of size 1 is ignored. While in Conv kernel,
+    # this tensor is considered as channels first and the output will be in contiguous format.
+    # To align the behavior of the Conv kernel, we set the output_stride in such case to be contiguous instead of channels last.
+    dynamic_shapes = not all(isinstance(i, int) for i in (output_size))
+    if (
+        dynamic_shapes or get_device_type(x) == "xpu"
+    ) and is_contiguous_storage_and_layout(x):
+        output_stride: StrideType = FlexibleLayout.contiguous_strides(output_size)
+    # Currently we don't support channel last for the situation that stride of input's batch dim is 0,
+    # eg. input_size = (1, 1280, 64, 64), but input_stride=(0, 1, 81920, 1280).
+    # So we use NCHW hear instead.
+    # Different with cpu, cpu conv always use channels_last for convolution when weight is prepacked,
+    # but xpu does not do the prepack, so the problem exposed here is only for xpu.
+    # TODO support channels_last for such zero stride input.
+    elif get_device_type(x) == "xpu" and x.get_stride()[0] == 0:
+        output_stride = FlexibleLayout.contiguous_strides(output_size)
+    else:
+        output_stride = make_channels_last_strides_for(output_size)
+
+    assert get_device_type(x) == get_device_type(weight)
+    assert get_device_type(x) in SUPPORTED_MKLDNN_DEVICES
+    inputs = [x]
+
+    if quantize_args is not None:
+        x_scale, x_zero_point, w_scale, w_zero_point = quantize_args
+        x_scale.realize()
+        x_zero_point.realize()
+        w_scale.realize()
+        w_zero_point.realize()
+        inputs = inputs + [x_scale, x_zero_point] + [weight] + [w_scale, w_zero_point]
+    else:
+        inputs += [weight]
+
+    if other is not None:
+        other = cls.require_stride_order(other, req_stride_order)
+        assert isinstance(other, TensorBox)
+        inputs += [other]
+
+    kernel_layout = FixedLayout(
+        x.get_device_or_error(),
+        x.get_dtype(),
+        convert_shape_to_inductor(output_size),
+        convert_shape_to_inductor(output_stride),
+    )
+    constant_args = [padding, stride, dilation, groups]
+    if transposed:
+        constant_args.insert(1, output_padding)
+
+    if bias is not None:
+        inputs.append(bias)
+    else:
+        constant_args.insert(0, bias)
+    return inputs, constant_args, kernel_layout, req_stride_order, other
+
+
+def _prepare_linear_fusion_create(
+    cls,
+    x: "TensorBox",
+    weight: "TensorBox",
+    bias: "TensorBox",
+    quantize_args: list["TensorBox"] | None = None,
+    other: Optional["TensorBox"] = None,
+    binary_sum: bool = False,
+):
+    """
+    This function is a helper function to prepare inputs, layout and constant args
+    for linear post-op fusion's create function. The function only supports the CPU device
+    since linear post-op fusion kernel is only supported on CPU right now.
+    """
+    x.realize()
+    weight.realize()
+    if bias is not None:
+        bias.realize()
+
+    *m, _ = x.get_size()
+    # The weight has been transposed during the qlinear weight prepack process.
+    # https://github.com/pytorch/pytorch/blob/4979f9c0d72490970e2019bb1d2284f83d93f76b/
+    # aten/src/ATen/native/quantized/cpu/qlinear_prepack.cpp#L291
+    _, oc = weight.get_size()
+    output_size = list(m) + [oc]
+    req_stride_order = list(reversed(range(len(x.get_size()))))
+
+    x = cls.require_stride_order(x, req_stride_order)
+    assert get_device_type(x) == get_device_type(weight)
+    assert get_device_type(x) in SUPPORTED_MKLDNN_DEVICES
+    inputs = [x]
+
+    if quantize_args is not None:
+        x_scale, x_zero_point, w_scale, w_zero_point = quantize_args
+        x_scale.realize()
+        x_zero_point.realize()
+        w_scale.realize()
+        w_zero_point.realize()
+        inputs = inputs + [x_scale, x_zero_point] + [weight] + [w_scale, w_zero_point]
+    else:
+        inputs += [weight]
+
+    if other is not None:
+        if binary_sum:
+            other = cls.require_stride_order(other, req_stride_order)
+        inputs = inputs + [other]
+
+    output_stride = FlexibleLayout.contiguous_strides(output_size)
+    kernel_layout = FixedLayout(
+        x.get_device(),
+        x.get_dtype(),
+        output_size,
+        output_stride,
+    )
+    constant_args: list[Any] = []
+
+    if bias is not None:
+        inputs.append(bias)
+    else:
+        constant_args.insert(0, bias)
+    return inputs, constant_args, kernel_layout, req_stride_order, other
+
+
+def _create_output_node(packed):
+    output_ir = MultiOutput(
+        packed.get_layout(),
+        packed,
+        [],
+    )
+    packed.layout = MultiOutputLayout(device=packed.get_device())
+    packed.outputs = [output_ir]
+    return output_ir
+
+
+class ConvolutionUnary(ExternKernelAlloc):
+    def __init__(
+        self,
+        layout,
+        inputs,
+        constant_args=(),
+    ) -> None:
+        self.device_type = get_device_type(inputs[0])
+        super().__init__(
+            layout,
+            inputs,
+            constant_args,
+            None,
+            op_overload=torch.ops.mkldnn._convolution_pointwise.default,
+            cpp_kernel_name=f"aoti_torch_{self.device_type}_mkldnn__convolution_pointwise",
+        )
+
+    def codegen(self, wrapper):
+        wrapper.include_extra_header(
+            f"torch/csrc/inductor/aoti_torch/c/shim_{self.device_type}.h"
+        )
+        super().codegen(wrapper)
+
+    @classmethod
+    def create(
+        cls,
+        x: "TensorBox",
+        weight: "TensorBox",
+        bias: "TensorBox",
+        padding_: list[int],
+        stride_: list[int],
+        dilation_: list[int],
+        groups: int,
+        attr,
+        scalars: list[Any] | None,
+        algorithm,
+    ):
+        (
+            inputs,
+            constant_args,
+            kernel_layout,
+            _,
+            _,
+        ) = _prepare_convolution_fusion_create(
+            cls, x, weight, bias, padding_, stride_, dilation_, groups
+        )
+        constant_args = constant_args + [
+            attr,
+            may_convert_to_optional(scalars),
+            algorithm,
+        ]
+        packed = ConvolutionUnary(
+            layout=kernel_layout,
+            inputs=inputs,
+            constant_args=constant_args,
+        )
+        return _create_output_node(packed)
+
+
+class ConvolutionBinary(ExternKernelAlloc):
+    def __init__(
+        self,
+        layout,
+        inputs,
+        constant_args=(),
+        cpp_constant_args=(),
+    ) -> None:
+        self.device_type = get_device_type(inputs[0])
+        super().__init__(
+            layout,
+            inputs,
+            constant_args,
+            None,
+            op_overload=torch.ops.mkldnn._convolution_pointwise.binary,
+            cpp_kernel_name=f"aoti_torch_{self.device_type}_mkldnn__convolution_pointwise_binary",
+        )
+        self.cpp_constant_args = cpp_constant_args
+
+    def codegen(self, wrapper):
+        wrapper.include_extra_header(
+            f"torch/csrc/inductor/aoti_torch/c/shim_{self.device_type}.h"
+        )
+        super().codegen(wrapper)
+
+    @classmethod
+    def create(
+        cls,
+        x: "TensorBox",
+        other: "TensorBox",
+        weight: "TensorBox",
+        bias: "TensorBox",
+        padding_: list[int],
+        stride_: list[int],
+        dilation_: list[int],
+        groups: int,
+        binary_attr: str,
+        binary_alpha: float | None,
+        unary_attr: str | None,
+        unary_scalars: list[Any] | None,
+        unary_algorithm: str | None,
+    ):
+        (
+            inputs,
+            constant_args,
+            kernel_layout,
+            req_stride_order,
+            _,
+        ) = _prepare_convolution_fusion_create(
+            cls, x, weight, bias, padding_, stride_, dilation_, groups
+        )
+        # pyrefly: ignore [bad-assignment]
+        other = cls.require_stride_order(other, req_stride_order)
+        inputs.insert(1, other)
+        constant_args = constant_args + [
+            binary_attr,
+            binary_alpha,
+            unary_attr,
+            may_convert_to_optional(unary_scalars),
+            unary_algorithm,
+        ]
+        packed = ConvolutionBinary(
+            layout=kernel_layout,
+            inputs=inputs,
+            constant_args=constant_args,
+        )
+        return _create_output_node(packed)
+
+
+class ConvolutionBinaryInplace(ExternKernelAlloc):
+    def __init__(
+        self,
+        kernel_layout,
+        inputs,
+        constant_args=(),
+    ) -> None:
+        # Due to constrain of op.call, other (Tensor&) should be at input[0]
+        self.device_type = get_device_type(inputs[0])
+        reordered_inputs = [inputs[1], inputs[0]] + inputs[2:]
+
+        super().__init__(
+            kernel_layout,
+            reordered_inputs,
+            constant_args,
+            None,
+            op_overload=torch.ops.mkldnn._convolution_pointwise_.binary,
+            cpp_kernel_name=f"aoti_torch_{self.device_type}_mkldnn__convolution_pointwise_binary_",
+        )
+
+        self.mutation_outputs = [
+            MutationOutput(NoneLayout(device=inputs[0].get_device()), inputs[0], self),
+            MutationOutput(NoneLayout(device=inputs[1].get_device()), inputs[1], self),
+        ]
+
+    def codegen(self, wrapper):
+        wrapper.include_extra_header(
+            f"torch/csrc/inductor/aoti_torch/c/shim_{self.device_type}.h"
+        )
+        super().codegen(wrapper)
+
+    def get_unbacked_symbol_defs(self) -> OrderedSet[sympy.Symbol]:
+        return OrderedSet()
+
+    @classmethod
+    def create(
+        cls,
+        x: "TensorBox",
+        other: "TensorBox",
+        weight: "TensorBox",
+        bias: "TensorBox",
+        padding_: list[int],
+        stride_: list[int],
+        dilation_: list[int],
+        groups: int,
+        binary_attr: str,
+        binary_alpha: float | None,
+        unary_attr: str | None,
+        unary_scalars: list[Any] | None,
+        unary_algorithm: str | None,
+    ):
+        (
+            inputs,
+            constant_args,
+            _,
+            req_stride_order,
+            _,
+        ) = _prepare_convolution_fusion_create(
+            cls, x, weight, bias, padding_, stride_, dilation_, groups
+        )
+        # pyrefly: ignore [bad-assignment]
+        other = cls.require_stride_order(other, req_stride_order)
+        inputs.insert(1, other)
+        constant_args = constant_args + [
+            binary_attr,
+            binary_alpha,
+            unary_attr,
+            may_convert_to_optional(unary_scalars),
+            unary_algorithm,
+        ]
+        packed = ConvolutionBinaryInplace(
+            kernel_layout=NoneLayout(device=inputs[1].get_device()),  # type: ignore[arg-type]
+            inputs=inputs,
+            constant_args=constant_args,
+        )
+        # This op mutates in place which means that the result is not the
+        # target but rather the input that is being mutated
+        # init reorders the inputs, so inputs[1] becomes packed.inputs[0]
+        return packed.inputs[0]
+
+
+class ConvolutionTransposeUnary(ExternKernelAlloc):
+    def __init__(
+        self,
+        layout,
+        inputs,
+        constant_args=(),
+    ) -> None:
+        self.device_type = get_device_type(inputs[0])
+        super().__init__(
+            layout,
+            inputs,
+            constant_args,
+            None,
+            op_overload=torch.ops.mkldnn._convolution_transpose_pointwise.default,
+            cpp_kernel_name=f"aoti_torch_{self.device_type}_mkldnn__convolution_transpose_pointwise",
+        )
+
+    def codegen(self, wrapper):
+        wrapper.include_extra_header(
+            f"torch/csrc/inductor/aoti_torch/c/shim_{self.device_type}.h"
+        )
+        super().codegen(wrapper)
+
+    @classmethod
+    def create(
+        cls,
+        x: "TensorBox",
+        weight: "TensorBox",
+        bias: "TensorBox",
+        padding_: list[int],
+        output_padding_: list[int],
+        stride_: list[int],
+        dilation_: list[int],
+        groups_: int,
+        attr,
+        scalars: list[Any] | None,
+        algorithm,
+    ):
+        transposed = True
+        (
+            inputs,
+            constant_args,
+            kernel_layout,
+            _,
+            _,
+        ) = _prepare_convolution_fusion_create(
+            cls,
+            x,
+            weight,
+            bias,
+            padding_,
+            stride_,
+            dilation_,
+            groups_,
+            transposed,
+            output_padding_,
+        )
+        constant_args = constant_args + [
+            attr,
+            may_convert_to_optional(scalars),
+            algorithm,
+        ]
+        packed = ConvolutionTransposeUnary(
+            layout=kernel_layout,
+            inputs=inputs,
+            constant_args=constant_args,
+        )
+        return _create_output_node(packed)
+
+
+class QConvPointWisePT2E(ExternKernelAlloc):
+    def __init__(
+        self,
+        layout,
+        inputs,
+        constant_args=(),
+    ) -> None:
+        """
+        if bias is not None
+            - inputs = [x, w, b, weight_scale, weight_zp]
+            - const_args is: [stride, padding, dilation, groups, x_scale, x_zp, o_scale, o_zp,
+              fp32_output, unary_attr, unary_scalars, unary_algorithm]
+        else
+            - inputs = [x, w, weight_scale, weight_zp]
+            - const_args is: [bias, stride, padding, dilation, groups, x_scale, x_zp, o_scale, o_zp,
+              fp32_output, unary_attr, unary_scalars, unary_algorithm]
+        """
+        self.device_type = get_device_type(inputs[0])
+        self.has_bias = len(inputs) == 5
+        super().__init__(
+            layout,
+            inputs,
+            constant_args,
+            None,
+            op_overload=torch.ops.onednn.qconv_pointwise.tensor,
+            cpp_kernel_name=f"aoti_torch_{self.device_type}__qconv_pointwise_tensor",
+        )
+
+    def codegen(self, wrapper):
+        wrapper.include_extra_header(
+            f"torch/csrc/inductor/aoti_torch/c/shim_{self.device_type}.h"
+        )
+        super().codegen(wrapper)
+        if isinstance(self.layout, Layout):
+            self.codegen_size_asserts(wrapper)
+
+    @classmethod
+    def create(
+        cls,
+        qx: "TensorBox",
+        x_scale: "TensorBox",
+        x_zero_point: "TensorBox",
+        qw: "TensorBox",  # qw
+        w_scale: "TensorBox",
+        w_zero_point,
+        bias: "TensorBox",
+        stride: list[int],
+        padding: list[int],
+        dilation: list[int],
+        groups: int,
+        output_scale: float,
+        output_zero_point: int,
+        output_dtype,
+        attr,
+        scalars,
+        algorithm,
+    ):
+        transposed = False
+        output_padding = None
+        (
+            inputs,
+            constant_args,
+            kernel_layout,
+            _,
+            _,
+        ) = _prepare_convolution_fusion_create(
+            cls,
+            qx,
+            qw,
+            bias,
+            padding,
+            stride,
+            dilation,
+            groups,
+            transposed,
+            output_padding,
+            [x_scale, x_zero_point, w_scale, w_zero_point],  # type: ignore[list-item]
+        )
+        # swap padding and stride to align with functional conv arg order
+        if bias is None:
+            constant_args[1], constant_args[2] = constant_args[2], constant_args[1]
+        else:
+            constant_args[0], constant_args[1] = constant_args[1], constant_args[0]
+
+        constant_args = constant_args + [
+            output_scale,
+            output_zero_point,
+            output_dtype,
+            attr,
+            may_convert_to_optional(scalars),
+            algorithm,
+        ]
+
+        assert output_dtype is not None
+        if output_dtype in [torch.float32, torch.bfloat16]:
+            # in _prepare_convolution_fusion_create, we use x.dtype (uint8) to create kernel_layout
+            # if we set output_dtype is not None, the output buf should be output_dtype instead of uint8.
+            kernel_layout.dtype = output_dtype
+
+        return QConvPointWisePT2E(
+            layout=kernel_layout,
+            inputs=inputs,
+            constant_args=constant_args,
+        )
+
+
+class QConvPointWiseBinaryPT2E(ExternKernelAlloc):
+    def __init__(
+        self,
+        layout,
+        inputs,
+        constant_args=(),
+    ) -> None:
+        """
+        Needs input/weight/output qparams
+        if bias is not None
+            - inputs = [x, x_scale, x_zp, w,  w_scale, w_zp, accum, b]
+            - const_args = [stride, padding, dilation, groups, o_scale, o_zp,
+            output_dtype, accum_scale, accum_zp, binary_attr, alpha, unary_attr, unary_scalars, unary_algorithm]
+        else
+            - inputs = [x, x_scale, x_zp, w,  w_scale, w_zp, accum]
+            - const_args [b, stride, padding, dilation, groups, o_scale, o_zp,
+             output_dtype, accum_scale, accum_zp, binary_attr, alpha, unary_attr, unary_scalars, unary_algorithm]
+        """
+        self.device_type = get_device_type(inputs[0])
+        self.has_bias = len(inputs) == 8
+        self.idx_for_inplace_sum = 6
+        super().__init__(
+            layout,
+            inputs,
+            constant_args,
+            None,
+            op_overload=torch.ops.onednn.qconv2d_pointwise.binary_tensor,
+            cpp_kernel_name=(
+                f"aoti_torch_{self.device_type}__qconv2d_pointwise_binary_tensor"
+            ),
+        )
+
+    def codegen(self, wrapper):
+        wrapper.include_extra_header(
+            f"torch/csrc/inductor/aoti_torch/c/shim_{self.device_type}.h"
+        )
+        super().codegen(wrapper)
+        if isinstance(self.layout, Layout):
+            self.codegen_size_asserts(wrapper)
+
+    def get_mutation_names(self) -> Sequence[str]:
+        return [self.input_name(self.idx_for_inplace_sum)]
+
+    def get_unbacked_symbol_defs(self) -> OrderedSet[sympy.Symbol]:
+        return OrderedSet()
+
+    @classmethod
+    def create(
+        cls,
+        qx: "TensorBox",
+        x_scale: "TensorBox",
+        x_zero_point: "TensorBox",
+        qw: "TensorBox",  # packed_weight
+        w_scale,
+        w_zero_point,
+        qaccum: "TensorBox",
+        bias: "TensorBox",
+        stride: list[int],
+        padding: list[int],
+        dilation: list[int],
+        groups: int,
+        output_scale: "TensorBox",
+        output_zero_point: "TensorBox",
+        output_dtype,
+        accum_scale,
+        accum_zero_point,
+        binary_attr,
+        alpha,
+        unary_attr,
+        unary_scalars,
+        unary_algorithm,
+    ):
+        transposed = False
+        output_padding = None
+        (
+            inputs,
+            constant_args,
+            _kernel_layout,
+            req_stride_order,
+            qaccum,
+        ) = _prepare_convolution_fusion_create(
+            cls,
+            qx,
+            qw,
+            bias,
+            padding,
+            stride,
+            dilation,
+            groups,
+            transposed,
+            output_padding,
+            [x_scale, x_zero_point, w_scale, w_zero_point],
+            qaccum,
+        )
+
+        # swap padding and stride to align with functional conv arg order
+        if bias is None:
+            constant_args[1], constant_args[2] = constant_args[2], constant_args[1]
+        else:
+            constant_args[0], constant_args[1] = constant_args[1], constant_args[0]
+
+        constant_args = constant_args + [
+            output_scale,
+            output_zero_point,
+            output_dtype,
+            accum_scale,
+            accum_zero_point,
+            binary_attr,
+            alpha,
+            unary_attr,
+            may_convert_to_optional(unary_scalars),
+            unary_algorithm,
+        ]
+
+        assert binary_attr == "sum", (
+            "For now, only post op sum is supported in QConvPointWiseBinaryPT2E."
+        )
+
+        V.graph.mark_buffer_mutated(qaccum.get_name())
+        packed = QConvPointWiseBinaryPT2E(
+            layout=NoneLayout(device=qaccum.get_device()),
+            inputs=inputs,
+            constant_args=constant_args,
+        )
+
+        # Return accum since it has been inplace changed.
+        return packed.inputs[packed.idx_for_inplace_sum]
+
+
+class MKLPackedLinear(ExternKernelAlloc):
+    def __init__(
+        self,
+        layout,
+        inputs,
+        constant_args=(),
+    ) -> None:
+        super().__init__(
+            layout,
+            inputs,
+            constant_args,
+            None,
+            op_overload=torch.ops.mkl._mkl_linear.default,
+        )
+
+    def codegen(self, wrapper):
+        wrapper.include_extra_header("torch/csrc/inductor/aoti_torch/c/shim_cpu.h")
+        super().codegen(wrapper)
+
+    @classmethod
+    def create(cls, x, packed_w, orig_w, B, batch_size):
+        x = cls.require_stride1(cls.realize_input(x))
+        orig_w = cls.require_stride1(cls.realize_input(orig_w))
+        *m, _ = x.get_size()
+        oc, _ = orig_w.get_size()
+        output_size = list(m) + [oc]
+        output_stride = FlexibleLayout.contiguous_strides(output_size)
+        inputs = [x, packed_w, orig_w]
+        constant_args = [batch_size]
+        if B is not None:
+            inputs += [B]
+        else:
+            constant_args.insert(0, None)
+
+        device = x.get_device()
+        assert device is not None
+        return MKLPackedLinear(
+            layout=FixedLayout(device, x.get_dtype(), output_size, output_stride),
+            inputs=inputs,
+            constant_args=constant_args,
+        )
+
+
+class LinearUnary(ExternKernelAlloc):
+    def __init__(
+        self,
+        layout,
+        inputs,
+        constant_args=(),
+    ) -> None:
+        self.device_type = get_device_type(inputs[0])
+        super().__init__(
+            layout,
+            inputs,
+            constant_args,
+            None,
+            op_overload=torch.ops.mkldnn._linear_pointwise.default,
+            cpp_kernel_name=f"aoti_torch_{self.device_type}__linear_pointwise",
+        )
+
+    def codegen(self, wrapper):
+        wrapper.include_extra_header(
+            f"torch/csrc/inductor/aoti_torch/c/shim_{self.device_type}.h"
+        )
+        super().codegen(wrapper)
+
+    @classmethod
+    def create(cls, x, w, B, attr, scalars, algorithm):
+        x = cls.require_contiguous(cls.realize_input(x))
+        w = cls.require_contiguous(cls.realize_input(w))
+
+        *m, _ic = x.get_size()
+        oc, _ic = w.get_size()
+        output_size = list(m) + [oc]
+        inputs = [x, w]
+        constant_args = [attr, scalars if scalars else [-1], algorithm]
+        if B is not None:
+            B = cls.require_contiguous(cls.realize_input(B))
+            inputs.append(B)
+        else:
+            constant_args.insert(0, None)
+
+        device = x.get_device()
+        assert device is not None
+
+        packed = LinearUnary(
+            layout=FixedLayout(
+                device=device,
+                dtype=x.get_dtype(),
+                size=output_size,
+            ),
+            inputs=inputs,
+            constant_args=constant_args,
+        )
+        return _create_output_node(packed)
+
+    def apply_constraint(self):
+        pass
+
+
+class LinearBinary(ExternKernelAlloc):
+    kernel = "torch.ops.mkldnn._linear_pointwise.binary"
+
+    def __init__(
+        self,
+        layout,
+        inputs,
+        constant_args=(),
+    ) -> None:
+        self.device_type = get_device_type(inputs[0])
+        super().__init__(
+            layout,
+            inputs,
+            constant_args,
+            None,
+            op_overload=torch.ops.mkldnn._linear_pointwise.binary,
+            cpp_kernel_name=f"aoti_torch_{self.device_type}__linear_pointwise_binary",
+        )
+
+    def codegen(self, wrapper):
+        wrapper.include_extra_header(
+            f"torch/csrc/inductor/aoti_torch/c/shim_{self.device_type}.h"
+        )
+        super().codegen(wrapper)
+
+    @classmethod
+    def create(cls, x, y, w, B, attr):
+        x = cls.require_contiguous(cls.realize_input(x))
+        y = cls.require_contiguous(cls.realize_input(y))
+        w = cls.require_contiguous(cls.realize_input(w))
+
+        *m, _ic = x.get_size()
+        oc, _ic = w.get_size()
+        output_size = list(m) + [oc]
+        inputs = [x, y, w]
+        constant_args = [attr]
+        if B is not None:
+            B = cls.require_contiguous(cls.realize_input(B))
+            inputs.append(B)
+        else:
+            constant_args.insert(0, B)
+
+        device = x.get_device()
+        assert device is not None
+        packed = LinearBinary(
+            layout=FixedLayout(
+                device=device,
+                dtype=x.get_dtype(),
+                size=output_size,
+            ),
+            inputs=inputs,
+            constant_args=constant_args,
+        )
+        return _create_output_node(packed)
+
+    def apply_constraint(self):
+        pass
+
+
+class QLinearPointwisePT2E(ExternKernelAlloc):
+    def __init__(
+        self,
+        layout,
+        inputs,
+        constant_args=(),
+        has_bias=True,
+    ) -> None:
+        """
+        if bias is not None
+            - inputs = [x, w, b, weight_scale, weight_zp]
+            - const_args is: [x_scale, x_zp, o_scale, o_zp,
+              fp32_output, unary_attr, unary_scalars, unary_algorithm]
+        else
+            - inputs = [x, w, weight_scale, weight_zp]
+            - const_args is: [bias, x_scale, x_zp, o_scale, o_zp,
+              fp32_output, unary_attr, unary_scalars, unary_algorithm]
+        """
+        self.device_type = get_device_type(inputs[0])
+        self.has_bias = has_bias
+        super().__init__(
+            layout,
+            inputs,
+            constant_args,
+            None,
+            op_overload=(torch.ops.onednn.qlinear_pointwise.tensor),
+            cpp_kernel_name=(
+                f"aoti_torch_{self.device_type}__qlinear_pointwise_tensor"
+            ),
+        )
+
+    def codegen(self, wrapper):
+        wrapper.include_extra_header(
+            f"torch/csrc/inductor/aoti_torch/c/shim_{self.device_type}.h"
+        )
+        super().codegen(wrapper)
+
+        if isinstance(self.layout, Layout):
+            self.codegen_size_asserts(wrapper)
+
+    @classmethod
+    def create(
+        cls,
+        qx: "TensorBox",
+        x_scale: "TensorBox",
+        x_zero_point: "TensorBox",
+        qw: "TensorBox",  # packed_weight
+        w_scale: "TensorBox",
+        w_zero_point: "TensorBox",
+        bias: "TensorBox",
+        output_scale: float,
+        output_zero_point: int,
+        output_dtype,
+        post_op_name,
+        post_op_args,
+        post_op_algorithm,
+    ):
+        (inputs, constant_args, kernel_layout, _, _) = _prepare_linear_fusion_create(
+            cls,
+            qx,
+            qw,
+            bias,
+            [x_scale, x_zero_point, w_scale, w_zero_point],
+        )
+
+        constant_args = constant_args + [
+            output_scale,
+            output_zero_point,
+            output_dtype,
+            post_op_name,
+            may_convert_to_optional(post_op_args),
+            post_op_algorithm,
+        ]
+
+        assert output_dtype is not None
+        if output_dtype in [torch.float32, torch.bfloat16]:
+            # in _prepare_linear_fusion_create, we use x.dtype (uint8) to create kernel_layout
+            # if we set fp32_output, the output buf should be dtype float32 instead of uint8.
+            kernel_layout.dtype = output_dtype
+
+        return QLinearPointwisePT2E(
+            layout=kernel_layout,
+            inputs=inputs,
+            constant_args=constant_args,
+            has_bias=(bias is not None),
+        )
+
+
+class QLinearPointwiseBinaryPT2E(ExternKernelAlloc):
+    def __init__(
+        self,
+        layout,
+        inputs,
+        constant_args=(),
+        has_bias=True,
+    ) -> None:
+        """
+        if bias is not None
+            - inputs = [x, w, x_scale, x_zp, weight_scale, weight_zp, x2, bias]
+            - const_args is: [o_scale, o_zp,
+              fp32_output, binary_attr, alpha, unary_attr, unary_scalars, unary_algorithm]
+        else
+            - inputs = [x, w, x_scale, x_zp, weight_scale, weight_zp, x2]
+            - const_args is: [bias, o_scale, o_zp,
+              fp32_output, binary_attr, alpha, unary_attr, unary_scalars, unary_algorithm]
+        """
+        self.device_type = get_device_type(inputs[0])
+        self.has_bias = has_bias
+        self.idx_for_inplace_sum = 6
+        super().__init__(
+            layout,
+            inputs,
+            constant_args,
+            None,
+            op_overload=(torch.ops.onednn.qlinear_pointwise.binary_tensor),
+            cpp_kernel_name=f"aoti_torch_{self.device_type}__qlinear_pointwise_binary_tensor",
+        )
+
+    def codegen(self, wrapper):
+        wrapper.include_extra_header(
+            f"torch/csrc/inductor/aoti_torch/c/shim_{self.device_type}.h"
+        )
+        super().codegen(wrapper)
+        if isinstance(self.layout, Layout):
+            self.codegen_size_asserts(wrapper)
+
+    def get_mutation_names(self) -> Sequence[str]:
+        binary_post_op = self.constant_args[-5]
+        if binary_post_op == "sum":
+            input = self.inputs[self.idx_for_inplace_sum]
+            assert isinstance(input, IRNode)
+            return [input.get_name()]
+        else:
+            return []
+
+    @classmethod
+    def create(
+        cls,
+        qx: "TensorBox",
+        x_scale: "TensorBox",
+        x_zero_point: "TensorBox",
+        qw: "TensorBox",  # packed_weight
+        w_scale: "TensorBox",
+        w_zero_point: "TensorBox",
+        other: "TensorBox",
+        bias: "TensorBox",
+        output_scale: float,
+        output_zero_point: int,
+        output_dtype,
+        other_scale,
+        other_zp,
+        binary_post_op,
+        binary_alpha,
+        unary_post_op,
+        unary_post_op_args,
+        unary_post_op_algorithm,
+    ):
+        (
+            inputs,
+            constant_args,
+            kernel_layout,
+            req_stride_order,
+            other,
+        ) = _prepare_linear_fusion_create(
+            cls,
+            qx,
+            qw,
+            bias,
+            [x_scale, x_zero_point, w_scale, w_zero_point],
+            other,
+            binary_post_op == "sum",
+        )
+
+        constant_args = constant_args + [
+            output_scale,
+            output_zero_point,
+            output_dtype,
+            other_scale,
+            other_zp,
+            binary_post_op,
+            binary_alpha,
+            unary_post_op,
+            may_convert_to_optional(unary_post_op_args),
+            unary_post_op_algorithm,
+        ]
+
+        if binary_post_op == "sum":
+            V.graph.mark_buffer_mutated(other.get_name())
+            packed = QLinearPointwiseBinaryPT2E(
+                layout=NoneLayout(device=other.get_device()),
+                inputs=inputs,
+                constant_args=constant_args,
+                has_bias=(bias is not None),
+            )
+            # Return other since it has been inplace changed.
+            return packed.inputs[packed.idx_for_inplace_sum]
+
+        assert output_dtype is not None
+        if output_dtype in [torch.float32, torch.bfloat16]:
+            # in _prepare_linear_fusion_create, we use x.dtype (uint8) to create kernel_layout
+            # if we set fp32_output, the output buf should be dtype float32 instead of uint8.
+            kernel_layout.dtype = output_dtype
+
+        return QLinearPointwiseBinaryPT2E(
+            layout=kernel_layout,
+            inputs=inputs,
+            constant_args=constant_args,
+            has_bias=(bias is not None),
+        )
+
+
+class MkldnnRnnLayer(ExternKernelAlloc):
+    def __init__(
+        self,
+        layout,
+        inputs,
+        constant_args=(),
+    ) -> None:
+        super().__init__(
+            layout,
+            inputs,
+            constant_args,
+            None,
+            op_overload=torch.ops.aten.mkldnn_rnn_layer.default,
+        )
+
+    @classmethod
+    def create(
+        cls,
+        x: "TensorBox",
+        w0: "TensorBox",
+        w1: "TensorBox",
+        w2: "TensorBox",
+        w3: "TensorBox",
+        hx: "TensorBox",
+        cx: "TensorBox",
+        reverse: bool,
+        batch_sizes: list[int],
+        mode: int,
+        hidden_size: int,
+        num_layers: int,
+        has_biases: bool,
+        bidirectional: bool,
+        batch_first: bool,
+        train: bool,
+    ):
+        # pyrefly: ignore [bad-assignment]
+        x = cls.require_stride1(cls.realize_input(x))
+        # If batch_first, x has been permuted in lstm before entering the mkldnn_rnn_layer.
+        # Make sure x is contiguous in batch_first case.
+        x.freeze_layout()
+        # pyrefly: ignore [bad-assignment]
+        w0 = cls.require_stride1(cls.realize_input(w0))
+        # pyrefly: ignore [bad-assignment]
+        w1 = cls.require_stride1(cls.realize_input(w1))
+        # pyrefly: ignore [bad-assignment]
+        w2 = cls.require_stride1(cls.realize_input(w2))
+        # pyrefly: ignore [bad-assignment]
+        w3 = cls.require_stride1(cls.realize_input(w3))
+        # pyrefly: ignore [bad-assignment]
+        hx = cls.require_stride1(cls.realize_input(hx))
+        hx.freeze_layout()
+        # pyrefly: ignore [bad-assignment]
+        cx = cls.require_stride1(cls.realize_input(cx))
+        cx.freeze_layout()
+
+        input_size = x.get_size()
+        assert len(input_size) == 3, "Expect lstm input to be 3D"
+        # batch_first is handled in the lstm OP. When entering
+        # rnn_layer here, we'll always have batch_first = False
+        seq_length, mini_batch, input_size = input_size
+        output_shape = [seq_length, mini_batch, hidden_size]
+
+        hy_shape = hx.get_size()
+        cy_shape = cx.get_size()
+
+        inputs = [x, w0, w1, w2, w3, hx, cx]
+        constant_args = [
+            reverse,
+            batch_sizes,
+            mode,
+            hidden_size,
+            num_layers,
+            has_biases,
+            bidirectional,
+            batch_first,
+            train,
+        ]
+
+        device = x.get_device()
+        assert device is not None
+        packed = MkldnnRnnLayer(
+            MultiOutputLayout(device=device),
+            inputs=inputs,
+            constant_args=constant_args,
+        )
+
+        def get_strides_of_lstm_output(output_shape, batch_first):
+            assert len(output_shape) == 3, "Expect output_shape to be 3D"
+            return FlexibleLayout.contiguous_strides(output_shape)
+
+        # C shim call requires all the outputs to be passed in, and thus the last
+        # dummy return value is added.
+        output_sizes = [output_shape, hy_shape, cy_shape, [1]]
+        output_strides = [
+            get_strides_of_lstm_output(output_shape, batch_first),
+            FlexibleLayout.contiguous_strides(hy_shape),
+            FlexibleLayout.contiguous_strides(cy_shape),
+            [1],
+        ]
+        output_ir = [
+            MultiOutput(
+                FixedLayout(
+                    x.get_device(),  # type: ignore[arg-type]
+                    x.get_dtype(),
+                    output_size,
+                    output_stride,
+                ),
+                packed,
+                [(tuple, i)],
+            )
+            for i, (output_size, output_stride) in enumerate(
+                zip(output_sizes, output_strides)
+            )
+        ]
+        packed.outputs = output_ir
+
+        return output_ir
+
+    def codegen(self, wrapper):
+        wrapper.include_extra_header("torch/csrc/inductor/aoti_torch/c/shim_cpu.h")
+        return super().codegen(wrapper)
+
+
+# Add this IR so that we can include shim_cpu.h for cpp_wrapper
+class WeightInt4PackMatmul(ExternKernelAlloc):
+    def __init__(
+        self,
+        layout,
+        inputs,
+        constant_args=(),
+    ) -> None:
+        """
+        inputs = [x, w, qGroupSize, qScalesAndZeros]
+        constant_args = ()
+        """
+        assert len(inputs) == 4
+        assert len(constant_args) == 0
+        super().__init__(
+            layout,
+            inputs,
+            constant_args,
+            None,
+            op_overload=(torch.ops.quantized.int4mm_packed_weight_cpu.default),
+            cpp_kernel_name=("aoti_torch_cpu__weight_int4pack_mm_cpu_tensor"),
+        )
+
+    def codegen(self, wrapper):
+        wrapper.include_extra_header("torch/csrc/inductor/aoti_torch/c/shim_cpu.h")
+        super().codegen(wrapper)
+
+        if isinstance(self.layout, Layout):
+            self.codegen_size_asserts(wrapper)
+
+    @classmethod
+    def create(
+        cls,
+        x: "TensorBox",
+        w: "TensorBox",
+        qGroupSize: "TensorBox",
+        qScalesAndZeros: "TensorBox",
+    ):
+        inputs = [x, w, qGroupSize, qScalesAndZeros]
+        *m, _ = x.get_size()
+        n, _ = w.get_size()
+        output_size = list(m) + [n]
+        output_stride = FlexibleLayout.contiguous_strides(output_size)
+        kernel_layout = FixedLayout(
+            x.get_device(),  # type: ignore[arg-type]
+            x.get_dtype(),
+            output_size,
+            output_stride,
+        )
+        return WeightInt4PackMatmul(
+            layout=kernel_layout,
+            inputs=inputs,
+        )

outputs/audit_venv/lib/python3.11/site-packages/torch/_inductor/mkldnn_lowerings.py

@@ -0,0 +1,1399 @@
+# mypy: allow-untyped-defs
+import functools
+
+import torch
+import torch.utils._pytree as pytree
+from torch._inductor.kernel.mm_common import mm_args
+
+from . import config, ir
+from .codegen.cpp_gemm_template import CppGemmTemplate
+from .codegen.cpp_grouped_gemm_template import CppGroupedGemmTemplate
+from .codegen.cpp_utils import create_epilogue_with_attr
+from .ir import TensorBox
+from .lowering import (
+    add,
+    add_needs_realized_inputs,
+    aten,
+    permute,
+    register_lowering,
+    to_dtype,
+    view,
+)
+from .select_algorithm import (
+    autotune_select_algorithm,
+    ChoiceCaller,
+    ExternKernelChoice,
+)
+from .utils import use_aten_gemm_kernels, use_cpp_gemm_template
+from .virtualized import ops, OpsValue, V
+
+
+def create_int8_compensation(
+    W_tensor: torch.Tensor,
+    packed_weight: ir.TensorBox,
+    x_scale: ir.TensorBox,
+    x_zp: ir.TensorBox,
+    w_scale: ir.TensorBox,
+) -> tuple[
+    bool,
+    ir.TensorBox,
+    ir.TensorBox | None,
+]:
+    x_w_scale: ir.TensorBox | None = None
+    use_int8_fast_compensation_path = all(
+        isinstance(item, ir.TensorBox)
+        and item.get_name() in V.graph.constants
+        and hasattr(item.data, "data")
+        and isinstance(item.data.data, ir.ConstantBuffer)
+        for item in [x_scale, x_zp, w_scale]
+    )
+    if use_int8_fast_compensation_path:
+        x_w_scale_tensor = (
+            V.graph.constants[x_scale.get_name()]
+            * V.graph.constants[w_scale.get_name()]
+        )
+        x_w_scale = V.graph.add_tensor_constant(
+            x_w_scale_tensor,
+            name=packed_weight.get_name() + "_x_w_compens",
+        )
+        weight_compens_tensor = torch.sum(W_tensor.to(torch.float), dim=0)
+        x_zp_tensor = V.graph.constants[x_zp.get_name()]
+        weight_compens_tensor = weight_compens_tensor * x_w_scale_tensor * x_zp_tensor
+        weight_compens = V.graph.add_tensor_constant(
+            weight_compens_tensor,
+            name=packed_weight.get_name() + "_BMatrixCompens",
+        )
+    else:
+        weight_compens_tensor = torch.sum(W_tensor.to(torch.float), dim=0)
+        weight_compens = V.graph.add_tensor_constant(
+            weight_compens_tensor,
+            name=packed_weight.get_name() + "_BMatrixCompens",
+        )
+    return (  # type: ignore[return-type]
+        use_int8_fast_compensation_path,
+        weight_compens,
+        x_w_scale,
+    )
+
+
+def codegen_int8_gemm_template_compensation(
+    use_int8_fast_compensation_path: bool,
+    input: OpsValue,
+    _weight_compo: OpsValue,
+    _x_scale: OpsValue | None,
+    _x_zp: OpsValue | None,
+    _w_scale: OpsValue | None,
+    _x_w_scale: OpsValue | None,
+) -> OpsValue:
+    if use_int8_fast_compensation_path:
+        temp = ops.sub(
+            ops.mul(
+                input,
+                _x_w_scale,
+            ),
+            _weight_compo,
+        )
+    else:
+        temp = ops.mul(
+            ops.mul(
+                input,
+                _x_scale,
+            ),
+            _w_scale,
+        )
+        # NOTE: We will apply compensation even if the x_zp is 0 for int8 quantization.
+        # That's because when torch.compile is invoked for dynamic quantization,
+        # x might coincidentally have such values that x_zp might be zero despite
+        # asymmetric quantization.
+        # Besides, if x_zp is dummy for int8 x, or if x is statically quantized,
+        # we'd still perform that redundant compute to avoid making the code messy
+        # because we discovered that redundant computation of compensation did not
+        # lead to performance degradation with the input shapes tested.
+        temp = ops.sub(
+            temp,
+            ops.mul(
+                ops.mul(
+                    ops.mul(
+                        _x_scale,
+                        _w_scale,
+                    ),
+                    _x_zp,
+                ),
+                _weight_compo,
+            ),
+        )
+    return temp
+
+
+def grouped_gemm_lowering(
+    x: TensorBox,
+    w: list[TensorBox],
+    b: list[TensorBox],
+    attr=None,
+    scalars=None,
+    algorithm=None,
+    layout=None,
+):
+    x_size = x.get_size()
+    if len(x_size) > 2:
+        # GEMM template needs 2D input, normalize input shape here
+        x = view(x, [-1, x_size[-1]])
+    num_gemm = len(w)
+
+    assert config.max_autotune or config.max_autotune_gemm
+    # pyrefly: ignore [bad-assignment]
+    b = [bias if bias is None else ir.ExternKernel.realize_input(bias) for bias in b]
+
+    choices: list[ChoiceCaller] = []
+    *_, layout, x, _ = mm_args(x, permute(w[0], [1, 0]), layout=layout)
+
+    kwargs = {
+        "has_bias": [bias is not None for bias in b],
+        "trans_w": True,
+        "epilogue_creator": None,
+        "act_mapping": dict.fromkeys(range(num_gemm), x),
+    }
+
+    input_nodes = [x, *w]
+    input_nodes.extend([bias for bias in b if bias is not None])
+
+    CppGroupedGemmTemplate.add_choices(
+        choices,
+        layout,
+        input_nodes,
+        **kwargs,  # type: ignore[arg-type]
+    )
+
+    assert len(choices) != 0
+    result, _ = autotune_select_algorithm(
+        "grouped_gemm",
+        choices,
+        input_nodes,
+        layout,
+    )
+    template_buf = result.data.data
+    return_bufs = [
+        ir.MultiOutput(layout, template_buf, [(list, gemm_idx)])
+        for gemm_idx in range(num_gemm)
+    ]
+    # pyrefly: ignore [bad-argument-type]
+    template_buf.layout = ir.MultiOutputLayout(device=input_nodes[0].get_device())
+    template_buf.outputs = return_bufs
+    return_tensors = [
+        ir.TensorBox.create(return_bufs[gemm_idx]) for gemm_idx in range(num_gemm)
+    ]
+    if len(x_size) > 2:
+        for gemm_idx in range(num_gemm):
+            return_tensors[gemm_idx] = view(
+                return_tensors[gemm_idx],  # type: ignore[arg-type]
+                (*x_size[:-1], return_tensors[gemm_idx].get_size()[-1]),
+            )
+    return return_tensors
+
+
+grouped_gemm_lowering._inductor_lowering_function = True  # type: ignore[attr-defined]
+
+
+def register_onednn_fusion_ops():
+    if torch._C._has_mkldnn:
+        from . import mkldnn_ir
+
+        aten_mkldnn_linear_unary = ExternKernelChoice(
+            torch.ops.mkldnn._linear_pointwise,
+            "mkldnn::_linear_pointwise",
+            has_out_variant=False,
+            kernel_creator=mkldnn_ir.LinearUnary.create,
+        )
+        aten_mkldnn_linear_binary = ExternKernelChoice(
+            torch.ops.mkldnn._linear_pointwise.binary,
+            "mkldnn::_linear_pointwise",
+            has_out_variant=False,
+            kernel_creator=mkldnn_ir.LinearBinary.create,
+        )
+        aten_mkldnn_qlinear_unary = ExternKernelChoice(
+            torch.ops.onednn.qlinear_pointwise,
+            "onednn::qlinear_pointwise",
+            has_out_variant=False,
+            kernel_creator=mkldnn_ir.QLinearPointwisePT2E.create,
+        )
+        aten_mkldnn_qlinear_binary = ExternKernelChoice(
+            torch.ops.onednn.qlinear_pointwise.binary,
+            "onednn::qlinear_pointwise",
+            has_out_variant=False,
+            kernel_creator=mkldnn_ir.QLinearPointwiseBinaryPT2E.create,
+        )
+        cpu_needs_realized_inputs: list[
+            torch._ops.OpOverload | torch._ops.OpOverloadPacket
+        ] = [
+            torch.ops.mkldnn._convolution_pointwise,
+            torch.ops.mkldnn._convolution_pointwise_,
+            torch.ops.mkldnn._convolution_transpose_pointwise,
+            torch.ops.mkldnn._linear_pointwise,
+            aten.mkldnn_rnn_layer.default,
+            torch.ops.onednn.qconv_pointwise,
+        ]
+
+        @register_lowering(torch.ops.mkldnn._convolution_pointwise)
+        def convolution_unary(
+            x: TensorBox,
+            weight: TensorBox,
+            bias: TensorBox,
+            padding,
+            stride,
+            dilation,
+            groups,
+            attr,
+            scalars,
+            algorithm,
+        ):
+            return TensorBox.create(
+                mkldnn_ir.ConvolutionUnary.create(
+                    x,
+                    weight,
+                    bias,
+                    padding,
+                    stride,
+                    dilation,
+                    groups,
+                    attr,
+                    scalars,
+                    algorithm,
+                )
+            )
+
+        @register_lowering(torch.ops.mkldnn._convolution_pointwise.binary)
+        def convolution_binary(
+            x: TensorBox,
+            other: TensorBox,
+            weight: TensorBox,
+            bias: TensorBox,
+            padding,
+            stride,
+            dilation,
+            groups,
+            binary_attr,
+            binary_alpha,
+            unary_attr,
+            unary_scalars,
+            unary_algorithm,
+        ):
+            return TensorBox.create(
+                mkldnn_ir.ConvolutionBinary.create(
+                    x,
+                    other,
+                    weight,
+                    bias,
+                    padding,
+                    stride,
+                    dilation,
+                    groups,
+                    binary_attr,
+                    binary_alpha,
+                    unary_attr,
+                    unary_scalars,
+                    unary_algorithm,
+                )
+            )
+
+        @register_lowering(torch.ops.mkldnn._convolution_pointwise_.binary)
+        def convolution_binary_inplace(
+            x: TensorBox,
+            other: TensorBox,
+            weight: TensorBox,
+            bias: TensorBox,
+            padding,
+            stride,
+            dilation,
+            groups,
+            binary_attr,
+            binary_alpha,
+            unary_attr,
+            unary_scalars,
+            unary_algorithm,
+        ):
+            return TensorBox.create(
+                mkldnn_ir.ConvolutionBinaryInplace.create(
+                    x,
+                    other,
+                    weight,
+                    bias,
+                    padding,
+                    stride,
+                    dilation,
+                    groups,
+                    binary_attr,
+                    binary_alpha,
+                    unary_attr,
+                    unary_scalars,
+                    unary_algorithm,
+                )
+            )
+
+        @register_lowering(torch.ops.mkldnn._linear_pointwise)
+        def linear_unary(
+            x: TensorBox,
+            w: TensorBox,
+            b: TensorBox,
+            attr,
+            scalars,
+            algorithm,
+            layout=None,
+        ):
+            x_size = x.get_size()
+            if len(x_size) > 2:
+                # GEMM template needs 2D input, normalize input shape here
+                x = view(x, [-1, x_size[-1]])
+            if b is not None:
+                b = ir.ExternKernel.realize_input(b)  # type: ignore[assignment]
+            choices: list[ChoiceCaller] = []
+            if config.max_autotune or config.max_autotune_gemm:
+                transposed_w = permute(w, [1, 0])
+                *_, layout, x, transposed_w = mm_args(x, transposed_w, layout=layout)
+                if use_cpp_gemm_template(layout, x, transposed_w):
+
+                    def epilogue_creator(buf):
+                        return create_epilogue_with_attr(
+                            buf, attr, scalars=scalars, algorithm=algorithm
+                        )
+
+                    kwargs = {
+                        "has_bias": b is not None,
+                        "trans_w": True,
+                        "epilogue_creator": (
+                            None if attr == "none" else epilogue_creator
+                        ),
+                    }
+                    if b is not None:
+                        kwargs["input_indices"] = [2, 0, 1]  # type: ignore[assignment]
+                    CppGemmTemplate.add_choices(
+                        choices,
+                        layout,
+                        [x, w] if b is None else [x, w, b],
+                        **kwargs,  # type: ignore[arg-type]
+                    )
+            if len(choices) == 0 or use_aten_gemm_kernels():
+                kwargs = dict(attr=attr, scalars=scalars, algorithm=algorithm)
+                if b is None:
+                    kwargs["B"] = None
+                choices.append(
+                    aten_mkldnn_linear_unary.bind(
+                        [x, w] if b is None else [x, w, b],
+                        layout,
+                        **kwargs,
+                    )
+                )
+            assert w.get_name() in V.graph.constants
+            input_gen_fns = {
+                1: lambda x: V.graph.constants[x.get_name()],
+            }
+            result, _ = autotune_select_algorithm(
+                "linear_unary",
+                choices,
+                [x, w] if b is None else [x, w, b],
+                layout,
+                input_gen_fns=input_gen_fns,
+            )
+            if len(x_size) > 2:
+                result = view(result, (*x_size[:-1], result.get_size()[-1]))
+            return result
+
+        @register_lowering(torch.ops.mkldnn._linear_pointwise.binary)
+        def linear_binary(
+            x: TensorBox, y: TensorBox, w: TensorBox, b: TensorBox, attr, layout=None
+        ):
+            x_size = x.get_size()
+            if len(x_size) > 2:
+                # GEMM template needs 2D input, normalize input shape here
+                x = view(x, [-1, x_size[-1]])
+            y_size = y.get_size()
+            if len(y_size) > 2:
+                y = view(y, [-1, y_size[-1]])
+            if b is not None:
+                b = ir.ExternKernel.realize_input(b)  # type: ignore[assignment]
+            choices: list[ChoiceCaller] = []
+            if config.max_autotune or config.max_autotune_gemm:
+                transposed_w = permute(w, [1, 0])
+                *_, layout, x, transposed_w, y = mm_args(
+                    x, transposed_w, y, layout=layout
+                )
+                if use_cpp_gemm_template(layout, x, transposed_w):
+
+                    def epilogue_creator(buf):
+                        return create_epilogue_with_attr(buf, attr, other=y)
+
+                    kwargs = {
+                        "has_bias": b is not None,
+                        "trans_w": True,
+                        "epilogue_creator": epilogue_creator,
+                    }
+
+                    # pyrefly: ignore [bad-typed-dict-key, unsupported-operation]
+                    kwargs["input_indices"] = [0, 2, 1] if b is None else [3, 0, 2, 1]
+                    CppGemmTemplate.add_choices(
+                        choices,
+                        layout,
+                        [x, y, w] if b is None else [x, y, w, b],
+                        **kwargs,  # type: ignore[arg-type]
+                    )
+            if len(choices) == 0 or use_aten_gemm_kernels():
+                kwargs = dict(attr=attr)
+                if b is None:
+                    kwargs["B"] = None
+                choices.append(
+                    aten_mkldnn_linear_binary.bind(
+                        [x, y, w] if b is None else [x, y, w, b],
+                        layout,
+                        **kwargs,
+                    )
+                )
+            assert w.get_name() in V.graph.constants
+            input_gen_fns = {
+                2: lambda x: V.graph.constants[x.get_name()],
+            }
+            result, _ = autotune_select_algorithm(
+                "linear_binary",
+                choices,
+                [x, y, w] if b is None else [x, y, w, b],
+                layout,
+                input_gen_fns=input_gen_fns,
+            )
+            if len(x_size) > 2:
+                result = view(result, (*x_size[:-1], result.get_size()[-1]))
+            return result
+
+        @register_lowering(torch.ops.mkldnn._convolution_transpose_pointwise)
+        def convolution_transpose_unary(
+            x: TensorBox,
+            weight: TensorBox,
+            bias: TensorBox,
+            padding,
+            output_padding,
+            stride,
+            dilation,
+            groups,
+            attr,
+            scalars,
+            algorithm,
+        ):
+            return TensorBox.create(
+                mkldnn_ir.ConvolutionTransposeUnary.create(
+                    x,
+                    weight,
+                    bias,
+                    padding,
+                    output_padding,
+                    stride,
+                    dilation,
+                    groups,
+                    attr,
+                    scalars,
+                    algorithm,
+                )
+            )
+
+        @register_lowering(aten.mkldnn_rnn_layer.default)
+        def mkldnn_rnn_layer(
+            x: TensorBox,
+            w0: TensorBox,
+            w1: TensorBox,
+            w2: TensorBox,
+            w3: TensorBox,
+            hx: TensorBox,
+            cx: TensorBox,
+            reverse: bool,
+            batch_sizes: list[int],
+            mode: int,
+            hidden_size: int,
+            num_layers: int,
+            has_biases: bool,
+            bidirectional: bool,
+            batch_first: bool,
+            train: bool,
+        ):
+            return pytree.tree_map(
+                TensorBox.create,
+                mkldnn_ir.MkldnnRnnLayer.create(
+                    x,
+                    w0,
+                    w1,
+                    w2,
+                    w3,
+                    hx,
+                    cx,
+                    reverse,
+                    batch_sizes,
+                    mode,
+                    hidden_size,
+                    num_layers,
+                    has_biases,
+                    bidirectional,
+                    batch_first,
+                    train,
+                ),
+            )
+
+        @register_lowering(torch.ops.onednn.qconv_pointwise, type_promotion_kind=None)
+        def qconvolution_unary(
+            x: TensorBox,
+            x_scale,
+            x_zp,
+            packed_weight: TensorBox,
+            w_scale: TensorBox,
+            w_zp,
+            bias: TensorBox,
+            stride,
+            padding,
+            dilation,
+            groups,
+            o_inv_scale,
+            o_zero_point,
+            output_dtype,
+            attr,
+            scalars,
+            algorithm,
+        ):
+            if not isinstance(x_scale, ir.TensorBox):
+                assert type(x_scale) is float
+                x_scale = V.graph.add_tensor_constant(
+                    torch.tensor(x_scale, dtype=torch.float32), name="x_scale"
+                )
+
+            if x_zp is None:
+                x_zp = V.graph.add_tensor_constant(
+                    torch.tensor(0, dtype=torch.int32), name="x_zp"
+                )
+            if not isinstance(x_zp, ir.TensorBox):
+                assert type(x_zp) is int
+                x_zp = V.graph.add_tensor_constant(
+                    torch.tensor(x_zp, dtype=torch.int32), name="x_zp"
+                )
+
+            if w_zp is None:
+                w_zp = V.graph.add_tensor_constant(
+                    torch.tensor(0, dtype=torch.int32), name="w_zp"
+                )
+
+            return TensorBox.create(
+                mkldnn_ir.QConvPointWisePT2E.create(
+                    x,
+                    x_scale,
+                    x_zp,
+                    packed_weight,
+                    w_scale,
+                    w_zp,
+                    bias,
+                    stride,
+                    padding,
+                    dilation,
+                    groups,
+                    o_inv_scale,
+                    o_zero_point,
+                    output_dtype,
+                    attr,
+                    scalars,
+                    algorithm,
+                )
+            )
+
+        @register_lowering(
+            torch.ops.onednn.qconv2d_pointwise.binary, type_promotion_kind=None
+        )
+        @register_lowering(
+            torch.ops.onednn.qconv2d_pointwise.binary_tensor, type_promotion_kind=None
+        )
+        def qconvolution_binary(
+            x: TensorBox,
+            x_scale,
+            x_zp,
+            packed_weight: TensorBox,
+            w_scale: TensorBox,
+            w_zp,
+            accum: TensorBox,
+            bias: TensorBox,
+            stride,
+            padding,
+            dilation,
+            groups,
+            o_inv_scale,
+            o_zero_point,
+            output_dtype,
+            accum_scale,
+            accum_zp,
+            binary_attr,
+            alpha,
+            unary_attr,
+            unary_scalars,
+            unary_algorithm,
+        ):
+            if not isinstance(x_scale, ir.TensorBox):
+                assert type(x_scale) is float
+                x_scale = V.graph.add_tensor_constant(
+                    torch.tensor(x_scale, dtype=torch.float32), name="x_scale"
+                )
+
+            if x_zp is None:
+                x_zp = V.graph.add_tensor_constant(
+                    torch.tensor(0, dtype=torch.int32), name="x_zp"
+                )
+            if not isinstance(x_zp, ir.TensorBox):
+                assert type(x_zp) is int
+                x_zp = V.graph.add_tensor_constant(
+                    torch.tensor(x_zp, dtype=torch.int32), name="x_zp"
+                )
+
+            if w_zp is None:
+                w_zp = V.graph.add_tensor_constant(
+                    torch.tensor(0, dtype=torch.int32), name="w_zp"
+                )
+
+            if (
+                binary_attr == "sum"
+                and output_dtype in [torch.float32, torch.bfloat16]
+                and accum.get_dtype() in [torch.float32, torch.bfloat16]
+                and accum.get_dtype() != output_dtype
+            ):
+                # For int8-mixed-bf16 quantization and inplace add,
+                # there is case when accum dtype is float32 but output dtype is bfloat16.
+                # Since the accum will be inplaced changed with post op sum,
+                # we will do accum dtype conversion here.
+                accum = to_dtype(accum, output_dtype)
+            return TensorBox.create(
+                mkldnn_ir.QConvPointWiseBinaryPT2E.create(
+                    x,
+                    x_scale,  # type: ignore[arg-type]
+                    x_zp,  # type: ignore[arg-type]
+                    packed_weight,
+                    w_scale,
+                    w_zp,
+                    accum,
+                    bias,
+                    stride,
+                    padding,
+                    dilation,
+                    groups,
+                    o_inv_scale,
+                    o_zero_point,
+                    output_dtype,
+                    accum_scale,
+                    accum_zp,
+                    binary_attr,
+                    alpha,
+                    unary_attr,
+                    unary_scalars,
+                    unary_algorithm,
+                )
+            )
+
+        @register_lowering(torch.ops.onednn.qlinear_pointwise, type_promotion_kind=None)
+        def qlinear_unary(
+            x: TensorBox,
+            x_scale,
+            x_zp,
+            packed_weight: TensorBox,
+            w_scale: TensorBox,
+            w_zp: TensorBox,
+            bias: TensorBox,
+            o_scale,
+            o_zero_point,
+            output_dtype,
+            attr,
+            scalars,
+            algorithm,
+            layout=None,
+        ):
+            assert packed_weight.get_dtype() in [torch.int8, torch.float8_e4m3fn], (
+                "Only int8 and e4m3fn weights are supported by oneDNN qlinear."
+            )
+            x_size = x.get_size()
+            if len(x_size) > 2:
+                # GEMM template needs 2D input, normalize input shape here
+                x = view(x, [-1, x_size[-1]])
+            if not isinstance(x_scale, ir.TensorBox):
+                assert type(x_scale) is float
+                x_scale = V.graph.add_tensor_constant(
+                    torch.tensor(x_scale, dtype=torch.float32), name="x_scale"
+                )
+            else:
+                x_scale.realize()
+                if all(dim == 1 for dim in x_scale.get_size()):
+                    # Corner-case discovered with LLaMA series.
+                    # If all outer dims of x_scale are 1, make it a 0D tensor.
+                    # Otherwise, epilogue creator will run into indexing issues.
+                    x_scale = view(x_scale, [])
+                assert len(x_scale.get_size()) in [0, 1], "x_scale must be 0D or 1D"
+
+            if x_zp is None:
+                # If x_zp is None, x is int8 quantized per-tensor and its scale is not reshaped,
+                # then the codegened code would segfault if we don't create a tensor for x_zp.
+                # It's safe to do so since x is a symmetrically quantized int8 tensor.
+                # Moreover, oneDNN qlinear API doesn't accept None value for zp
+                x_zp = V.graph.add_tensor_constant(
+                    torch.tensor(0, dtype=torch.int32), name="x_zp"
+                )
+            if not isinstance(x_zp, ir.TensorBox):
+                assert type(x_zp) is int
+                x_zp = V.graph.add_tensor_constant(
+                    torch.tensor(x_zp, dtype=torch.int32), name="x_zp"
+                )
+            else:
+                x_zp.realize()
+
+            assert x_zp.get_numel() == 1, "x_zp is incompatible with oneDNN qlinear"
+
+            # When channels less than 8, w_scale/w_zp is Pointwise instead of ConstantBuffer
+            # Refer to
+            # https://github.com/pytorch/pytorch/blob/f353d17755ed23b02924c962a86ff99a3405fe10/torch/_inductor/graph.py#L570-L577  # noqa: B950
+            if w_zp is None:
+                # If w_zp is None, then it's a dummy tensor created to denote the
+                # absence of a zero point, and thus w is int8 symmetrically quantized.
+                # Moreover, oneDNN qlinear API doesn't accept None value for zp
+
+                w_zp = V.graph.add_tensor_constant(
+                    torch.tensor(0, dtype=torch.int32), name="w_zp"
+                )
+            w_scale.realize()
+            w_zp.realize()
+            if w_zp.get_dtype() != torch.int32 and isinstance(
+                ir.InputsKernel.unwrap_storage_for_input(w_zp),
+                ir.ConstantBuffer,
+            ):
+                # W_zp might be a ConstantBuffer with int64, convert it to int32
+                w_zp_tensor = V.graph.constants[w_zp.get_name()].to(torch.int32)
+                w_zp = V.graph.add_tensor_constant(  # type: ignore[assignment]
+                    torch.tensor(w_zp_tensor, dtype=torch.int32), name=w_zp.get_name()
+                )
+
+            bias_dtype = None if bias is None else bias.get_dtype()
+            choices: list[ChoiceCaller] = []
+
+            if config.max_autotune or config.max_autotune_gemm:
+                *_, layout, x, packed_weight = mm_args(
+                    x, packed_weight, layout=layout, out_dtype=output_dtype
+                )
+
+                if (
+                    # GEMM template currently only supports symmetrically quantized weights
+                    isinstance(
+                        ir.InputsKernel.unwrap_storage_for_input(w_zp),
+                        ir.ConstantBuffer,
+                    )
+                    and torch.equal(
+                        torch.zeros_like(V.graph.constants[w_zp.get_name()]),
+                        V.graph.constants[w_zp.get_name()],
+                    )
+                ) and use_cpp_gemm_template(layout, x, packed_weight):
+                    W_tensor = V.graph.constants[packed_weight.get_name()].to_dense()
+
+                    (
+                        use_int8_fast_compensation_path,
+                        weight_compens,
+                        x_w_scale,
+                    ) = create_int8_compensation(
+                        W_tensor,
+                        packed_weight,
+                        x_scale,
+                        x_zp,
+                        w_scale,
+                    )
+
+                    def epilogue_creator(input_buffer):
+                        # Epilogue to convert from s32 to f32 for u8s8f32
+                        assert output_dtype in [
+                            torch.float32,
+                            torch.bfloat16,
+                            torch.uint8,
+                            torch.int8,
+                        ]
+                        input_loader = input_buffer.make_loader()
+                        weight_compens_loader = weight_compens.make_loader()
+                        x_w_scale_loader = None
+                        if use_int8_fast_compensation_path:
+                            assert x_w_scale is not None
+                            x_w_scale_loader = x_w_scale.make_loader()
+                        x_scale_loader = x_scale.make_loader()
+                        w_scale_loader = w_scale.make_loader()
+                        x_zp_loader = x_zp.make_loader()
+                        nonlocal bias
+                        bias_loader = None
+                        if bias is not None:
+                            bias_loader = bias.make_loader()
+
+                        def inner_fn(index):
+                            nonlocal bias
+                            input = input_loader(index)
+                            # MicroKernel Output is with int32
+                            # cvt to FP32 before doing compensation
+                            input = ops.to_dtype(input, torch.float32)
+                            weight_compens_index = (index[-1],)
+
+                            _x_scale = None
+                            _x_zp = None
+                            _w_scale = None
+                            if not use_int8_fast_compensation_path:
+                                _x_scale = x_scale_loader(())
+                                _x_zp = x_zp_loader(())
+                                _w_scale = w_scale_loader(weight_compens_index)
+                            _weight_compo = weight_compens_loader(weight_compens_index)
+                            _x_w_scale = None
+                            if use_int8_fast_compensation_path:
+                                assert x_w_scale_loader is not None
+                                _x_w_scale = x_w_scale_loader(weight_compens_index)
+                            # Step 1: Compute s8s8->s32 or u8s8->s32 GEMM & then apply compensation
+                            temp = codegen_int8_gemm_template_compensation(
+                                use_int8_fast_compensation_path,
+                                input,
+                                _weight_compo,
+                                _x_scale,
+                                _x_zp,
+                                _w_scale,
+                                _x_w_scale,
+                            )
+                            # Step 2: add Bias if applicable
+                            if bias is not None:
+                                # pyrefly: ignore [not-callable]
+                                _bias = bias_loader(weight_compens_index)
+                                nonlocal bias_dtype
+                                assert bias_dtype in [torch.float32, torch.bfloat16]
+                                if bias_dtype == torch.bfloat16:
+                                    _bias = ops.to_dtype(_bias, torch.float32)
+                                temp = ops.add(temp, _bias)
+
+                            return temp
+
+                        output_buf = ir.Pointwise(
+                            device=input_buffer.get_device(),
+                            dtype=torch.float32,  # Hardcode to FP32 for u8s8f32 & s8s8f32
+                            inner_fn=inner_fn,
+                            ranges=input_buffer.get_size(),
+                        )
+
+                        # Step 3: Doing the unary post op fusion
+                        if attr != "none":
+                            output_buf = create_epilogue_with_attr(
+                                output_buf, attr, scalars=scalars, algorithm=algorithm
+                            )
+
+                        # Step 4: Cast output to Target Dtype
+                        if output_dtype == torch.bfloat16:
+                            output_cast_loader = output_buf.make_loader()
+
+                            def inner_fn_cast_output_to_bf16(index):
+                                input = output_cast_loader(index)
+                                return ops.to_dtype(input, output_dtype)
+
+                            output_buf = ir.Pointwise(
+                                device=output_buf.get_device_or_error(),
+                                dtype=output_dtype,
+                                inner_fn=inner_fn_cast_output_to_bf16,
+                                ranges=output_buf.get_size(),
+                            )
+                        elif output_dtype in [torch.uint8, torch.int8]:
+                            from .lowering import _create_constants
+
+                            requant_input_loader = output_buf.make_loader()
+
+                            def inner_fn_requant(index, scale, zero_point):
+                                input = requant_input_loader(index)
+                                inv_scale, zero_point = _create_constants(
+                                    1.0 / scale, zero_point, dtype=torch.float32
+                                )
+                                val = ops.round(input * inv_scale) + zero_point
+                                if output_dtype == torch.uint8:
+                                    qmin, qmax = _create_constants(
+                                        0, 255, dtype=torch.float32
+                                    )
+                                else:
+                                    qmin, qmax = _create_constants(
+                                        -128, 127, dtype=torch.float32
+                                    )
+                                clamped = ops.minimum(ops.maximum(val, qmin), qmax)
+                                return ops.to_dtype(clamped, output_dtype)
+
+                            output_buf = ir.Pointwise(
+                                device=output_buf.get_device_or_error(),
+                                dtype=output_dtype,
+                                inner_fn=functools.partial(
+                                    inner_fn_requant,
+                                    scale=float(o_scale),
+                                    zero_point=int(o_zero_point),
+                                ),
+                                ranges=output_buf.get_size(),
+                            )
+
+                        return output_buf
+
+                    assert x.get_dtype() in [torch.uint8, torch.int8]
+                    CppGemmTemplate.add_choices(
+                        choices,
+                        layout,
+                        [x, x_scale, x_zp, packed_weight, w_scale, w_zp]
+                        if bias is None
+                        else [x, x_scale, x_zp, packed_weight, w_scale, w_zp, bias],
+                        has_bias=bias is not None,
+                        epilogue_creator=epilogue_creator,
+                        input_indices=[0, 3, 1, 2, 4, 5]
+                        if bias is None
+                        else [6, 0, 3, 1, 2, 4, 5],
+                    )
+            if len(choices) == 0 or use_aten_gemm_kernels():
+                kwargs = dict(
+                    output_scale=o_scale,
+                    output_zero_point=o_zero_point,
+                    output_dtype=output_dtype,
+                    post_op_name=attr,
+                    post_op_args=scalars,
+                    post_op_algorithm=algorithm,
+                )
+                if bias is None:
+                    kwargs["bias"] = None
+                choices.append(
+                    aten_mkldnn_qlinear_unary.bind(
+                        (x, x_scale, x_zp, packed_weight, w_scale, w_zp)
+                        if bias is None
+                        else (x, x_scale, x_zp, packed_weight, w_scale, w_zp, bias),
+                        layout,
+                        **kwargs,
+                    )
+                )
+            assert packed_weight.get_name() in V.graph.constants
+            input_gen_fns = {
+                3: lambda x: V.graph.constants[x.get_name()],  # packed weight
+                4: lambda x: V.graph.constants[x.get_name()],  # weight scale
+                5: lambda x: V.graph.constants[x.get_name()],  # weight zp
+                6: lambda x: V.graph.constants[x.get_name()],  # bias
+            }
+            if isinstance(
+                ir.InputsKernel.unwrap_storage_for_input(x_scale),
+                ir.ConstantBuffer,
+            ):
+                # x is statically quantized
+                input_gen_fns[1] = lambda x: V.graph.constants[x.get_name()]
+            if isinstance(
+                ir.InputsKernel.unwrap_storage_for_input(x_zp),
+                ir.ConstantBuffer,
+            ):
+                input_gen_fns[2] = lambda x: V.graph.constants[x.get_name()]
+
+            result, _ = autotune_select_algorithm(
+                "qlinear_unary",
+                choices,
+                [x, x_scale, x_zp, packed_weight, w_scale, w_zp]
+                if bias is None
+                else [x, x_scale, x_zp, packed_weight, w_scale, w_zp, bias],
+                layout,
+                input_gen_fns=input_gen_fns,
+            )
+            if len(x_size) > 2:
+                result = view(result, (*x_size[:-1], result.get_size()[-1]))
+            return result
+
+        @register_lowering(
+            torch.ops.onednn.qlinear_pointwise.binary, type_promotion_kind=None
+        )
+        @register_lowering(
+            torch.ops.onednn.qlinear_pointwise.binary_tensor, type_promotion_kind=None
+        )
+        def qlinear_binary(
+            x: TensorBox,
+            x_scale,
+            x_zp,
+            packed_weight: TensorBox,
+            w_scale: TensorBox,
+            w_zp: TensorBox,
+            x2: TensorBox,
+            bias: TensorBox,
+            o_scale,
+            o_zero_point,
+            output_dtype,
+            x2_scale,
+            x2_zp,
+            binary_attr,
+            alpha,
+            unary_attr,
+            unary_scalars,
+            unary_algorithm,
+            layout=None,
+        ):
+            x_size = x.get_size()
+            x2_size = x2.get_size()
+            assert len(x_size) == len(x2_size)
+            if len(x_size) > 2 and binary_attr in ["add", "sum"]:
+                # GEMM template needs 2D input, normalize input shape here
+                x = view(x, [-1, x_size[-1]])
+                x2 = view(x2, [-1, x2_size[-1]])
+            if not isinstance(x_scale, ir.TensorBox):
+                assert type(x_scale) is float
+                x_scale = V.graph.add_tensor_constant(
+                    torch.tensor(x_scale, dtype=torch.float32), name="x_scale"
+                )
+            else:
+                x_scale.realize()
+                if all(dim == 1 for dim in x_scale.get_size()):
+                    # Corner-case discovered with LLaMA series.
+                    # If all outer dims of x_scale are 1, make it a 0D tensor.
+                    # Otherwise, epilogue creator will run into indexing issues.
+                    x_scale = view(x_scale, [])
+                assert len(x_scale.get_size()) in [0, 1], "x_scale must be 0D or 1D"
+
+            if x_zp is None:
+                x_zp = V.graph.add_tensor_constant(
+                    torch.tensor(0, dtype=torch.int32), name="x_zp"
+                )
+
+            if w_zp is None:
+                w_zp = V.graph.add_tensor_constant(
+                    torch.tensor(0, dtype=torch.int32), name="w_zp"
+                )
+
+            if not isinstance(x_zp, ir.TensorBox):
+                assert type(x_zp) is int
+                x_zp = V.graph.add_tensor_constant(
+                    torch.tensor(x_zp, dtype=torch.int32), name="x_zp"
+                )
+            else:
+                x_zp.realize()
+
+            # When channels less than 8, w_scale/w_zp is Pointwise instead of ConstantBuffer
+            # Refer to
+            # https://github.com/pytorch/pytorch/blob/f353d17755ed23b02924c962a86ff99a3405fe10/torch/_inductor/graph.py#L570-L577  # noqa: B950
+            w_scale.realize()
+            w_zp.realize()
+            if w_zp.get_dtype() != torch.int32 and isinstance(
+                ir.InputsKernel.unwrap_storage_for_input(w_zp),
+                ir.ConstantBuffer,
+            ):
+                w_zp_tensor = V.graph.constants[w_zp.get_name()].to(torch.int32)
+                w_zp = V.graph.add_tensor_constant(  # type: ignore[assignment]
+                    torch.tensor(w_zp_tensor, dtype=torch.int32), name=w_zp.get_name()
+                )
+            if binary_attr == "sum":
+                if output_dtype in [
+                    torch.float32,
+                    torch.bfloat16,
+                ] and x2.get_dtype() in [torch.float32, torch.bfloat16]:
+                    if x2.get_dtype() != output_dtype:
+                        # For int8-mixed-bf16 quantization and inplace add,
+                        # there is case when accum dtype is float32 but output dtype is bfloat16.
+                        # Since the accum will be inplaced changed with post op sum,
+                        # we will do accum dtype conversion here.
+                        x2 = to_dtype(x2, output_dtype)
+                else:
+                    assert x2.get_dtype() == output_dtype, (
+                        "dtype of accum for qlinear post op sum should be the same as output"
+                    )
+            x2_dtype = x2.get_dtype()
+            bias_dtype = bias.get_dtype() if bias is not None else None
+            choices: list[ChoiceCaller] = []
+            if (config.max_autotune or config.max_autotune_gemm) and binary_attr in [
+                "add",
+                "sum",
+            ]:
+                *_, layout, x, packed_weight, x2 = mm_args(
+                    x, packed_weight, x2, layout=layout, out_dtype=output_dtype
+                )
+                if (
+                    isinstance(
+                        ir.InputsKernel.unwrap_storage_for_input(x_zp),
+                        ir.ConstantBuffer,
+                    )
+                    and len(x_zp.get_layout().size) == 0  # Per tensor quant of act
+                    and isinstance(
+                        ir.InputsKernel.unwrap_storage_for_input(w_zp),
+                        ir.ConstantBuffer,
+                    )
+                    and torch.equal(
+                        torch.zeros_like(V.graph.constants[w_zp.get_name()]),
+                        V.graph.constants[w_zp.get_name()],
+                    )  # We only compensate MatrixB and assume B_zp is 0 to avoid the compensation of MatrixA
+                    and use_cpp_gemm_template(layout, x, packed_weight)
+                ):
+                    W_tensor = V.graph.constants[packed_weight.get_name()]
+                    W_tensor = W_tensor.to_dense()
+                    (
+                        use_int8_fast_compensation_path,
+                        weight_compens,
+                        x_w_scale,
+                    ) = create_int8_compensation(
+                        W_tensor,
+                        packed_weight,
+                        x_scale,
+                        x_zp,
+                        w_scale,
+                    )
+
+                    def epilogue_creator(input_buffer):
+                        # Epilogue to convert from s32 to f32 for u8s8f32
+                        assert output_dtype in [
+                            torch.float32,
+                            torch.bfloat16,
+                            torch.uint8,
+                            torch.int8,
+                        ]
+
+                        input_loader = input_buffer.make_loader()
+                        x2_loader = x2.make_loader()
+                        weight_compens_loader = weight_compens.make_loader()
+                        x_w_scale_loader = None
+                        if use_int8_fast_compensation_path:
+                            assert x_w_scale is not None
+                            x_w_scale_loader = x_w_scale.make_loader()
+                        x_scale_loader = x_scale.make_loader()
+                        w_scale_loader = w_scale.make_loader()
+                        x_zp_loader = x_zp.make_loader()
+                        nonlocal bias
+                        bias_loader = None
+                        if bias is not None:
+                            bias_loader = bias.make_loader()
+
+                        def inner_fn(index):
+                            nonlocal bias
+                            input = input_loader(index)
+                            _x2 = x2_loader(index)
+                            _x_scale = None
+                            _x_zp = None
+                            _w_scale = None
+                            weight_compens_index = (index[-1],)
+                            if not use_int8_fast_compensation_path:
+                                _x_scale = x_scale_loader(())
+                                _x_zp = x_zp_loader(())
+                                _w_scale = w_scale_loader(weight_compens_index)
+                            # MicroKernel Output is with int32: cvt to FP32 before doing compensation
+                            input = ops.to_dtype(input, torch.float32)
+                            _weight_compo = weight_compens_loader(weight_compens_index)
+                            _x_w_scale = None
+                            if use_int8_fast_compensation_path:
+                                assert x_w_scale_loader is not None
+                                _x_w_scale = x_w_scale_loader(weight_compens_index)
+                            # Step 1: Doing compensation to cvt fp32
+                            temp = codegen_int8_gemm_template_compensation(
+                                use_int8_fast_compensation_path,
+                                input,
+                                _weight_compo,
+                                _x_scale,
+                                _x_zp,
+                                _w_scale,
+                                _x_w_scale,
+                            )
+                            # Step 2: add Bias if applicable
+                            if bias is not None:
+                                # pyrefly: ignore [not-callable]
+                                _bias = bias_loader(weight_compens_index)
+                                nonlocal bias_dtype
+                                assert bias_dtype in [torch.float32, torch.bfloat16]
+                                if bias_dtype == torch.bfloat16:
+                                    _bias = ops.to_dtype(_bias, torch.float32)
+                                temp = ops.add(temp, _bias)
+
+                            # Step 3: Binary add
+                            nonlocal x2_dtype
+                            assert x2_dtype in [torch.float32, torch.bfloat16]
+                            if x2_dtype == torch.bfloat16:
+                                _x2 = ops.to_dtype(_x2, torch.float32)
+                            temp = ops.add(temp, _x2)
+
+                            return temp
+
+                        output_buf = ir.Pointwise(
+                            device=input_buffer.get_device(),
+                            dtype=torch.float32,  # Hardcode to FP32 for u8s8f32
+                            inner_fn=inner_fn,
+                            ranges=input_buffer.get_size(),
+                        )
+
+                        # Step 4: Unary post op if has
+                        if unary_attr != "none":
+                            output_buf = create_epilogue_with_attr(
+                                output_buf,
+                                unary_attr,
+                                scalars=unary_scalars,
+                                algorithm=unary_algorithm,
+                            )
+
+                        # Step 5: Cast output to Target Dtype
+                        if output_dtype == torch.bfloat16:
+                            output_cast_loader = output_buf.make_loader()
+
+                            def inner_fn_cast_output_to_bf16(index):
+                                input = output_cast_loader(index)
+                                return ops.to_dtype(input, output_dtype)
+
+                            output_buf = ir.Pointwise(
+                                device=output_buf.get_device_or_error(),
+                                dtype=output_dtype,
+                                inner_fn=inner_fn_cast_output_to_bf16,
+                                ranges=output_buf.get_size(),
+                            )
+                        elif output_dtype in [torch.uint8, torch.int8]:
+                            from .lowering import _create_constants
+
+                            requant_input_loader = output_buf.make_loader()
+
+                            def inner_fn_requant(index, scale, zero_point):
+                                input = requant_input_loader(index)
+                                inv_scale, zero_point = _create_constants(
+                                    1.0 / scale, zero_point, dtype=torch.float32
+                                )
+                                val = ops.round(input * inv_scale) + zero_point
+                                if output_dtype == torch.uint8:
+                                    qmin, qmax = _create_constants(
+                                        0, 255, dtype=torch.float32
+                                    )
+                                else:
+                                    qmin, qmax = _create_constants(
+                                        -128, 127, dtype=torch.float32
+                                    )
+                                clamped = ops.minimum(ops.maximum(val, qmin), qmax)
+                                return ops.to_dtype(clamped, torch.uint8)
+
+                            output_buf = ir.Pointwise(
+                                device=output_buf.get_device_or_error(),
+                                dtype=torch.uint8,
+                                inner_fn=functools.partial(
+                                    inner_fn_requant,
+                                    scale=float(o_scale),
+                                    zero_point=int(o_zero_point),
+                                ),
+                                ranges=output_buf.get_size(),
+                            )
+
+                        return output_buf
+
+                    CppGemmTemplate.add_choices(
+                        choices,
+                        layout,
+                        [x, x_scale, x_zp, packed_weight, w_scale, w_zp, x2]
+                        if bias is None
+                        else [x, x_scale, x_zp, packed_weight, w_scale, w_zp, x2, bias],
+                        has_bias=bias is not None,
+                        epilogue_creator=epilogue_creator,
+                        # Reorder bias and x2
+                        input_indices=[0, 3, 1, 2, 4, 5, 6]
+                        if bias is None
+                        else [7, 0, 3, 1, 2, 4, 5, 6],
+                    )
+
+            if len(choices) == 0 or use_aten_gemm_kernels():
+                kwargs = dict(
+                    output_scale=o_scale,
+                    output_zero_point=o_zero_point,
+                    output_dtype=output_dtype,
+                    other_scale=x2_scale,
+                    other_zp=x2_zp,
+                    binary_post_op=binary_attr,
+                    binary_alpha=alpha,
+                    unary_post_op=unary_attr,
+                    unary_post_op_args=unary_scalars,
+                    unary_post_op_algorithm=unary_algorithm,
+                )
+                if bias is None:
+                    kwargs["bias"] = None
+                choices.append(
+                    aten_mkldnn_qlinear_binary.bind(
+                        (x, x_scale, x_zp, packed_weight, w_scale, w_zp, x2)
+                        if bias is None
+                        else (x, x_scale, x_zp, packed_weight, w_scale, w_zp, x2, bias),
+                        layout,
+                        **kwargs,
+                    )
+                )
+            assert packed_weight.get_name() in V.graph.constants
+            input_gen_fns = {
+                3: lambda x: V.graph.constants[x.get_name()],
+                4: lambda x: V.graph.constants[x.get_name()],
+                5: lambda x: V.graph.constants[x.get_name()],
+            }
+            if bias is not None:
+                input_gen_fns[7] = lambda x: V.graph.constants[x.get_name()]  # For bias
+            result, _ = autotune_select_algorithm(
+                "qlinear_binary",
+                choices,
+                [x, x_scale, x_zp, packed_weight, w_scale, w_zp, x2]
+                if bias is None
+                else [x, x_scale, x_zp, packed_weight, w_scale, w_zp, x2, bias],
+                layout,
+                input_gen_fns=input_gen_fns,
+            )
+            if (
+                isinstance(result.data.data, ir.CppTemplateBuffer)
+                and binary_attr == "sum"
+                and result.data.data.layout == x2.get_layout()
+            ):
+                # In this case, since x2 is inplace updated when binary_attr is "sum"
+                # we update the layout of result to view of x2
+                result = ir.TensorBox.create(
+                    ir.CppTemplateBuffer(
+                        layout=ir.NonOwningLayout(
+                            ir.ReinterpretView(data=x2, layout=x2.get_layout())
+                        ),
+                        inputs=result.data.data.inputs,  # type: ignore[arg-type]
+                        make_kernel_render=result.data.data.make_kernel_render,  # type: ignore[arg-type]
+                        template=result.data.data.template,
+                        choice=result.data.data.choice,
+                    )
+                )
+            if len(x_size) > 2 and binary_attr in ["add", "sum"]:
+                result = view(result, (*x_size[:-1], result.get_size()[-1]))  # type: ignore[arg-type]
+            return result
+
+        if torch._C.has_mkl:
+            aten_mkl_linear = ExternKernelChoice(
+                torch.ops.mkl._mkl_linear,
+                "mkl::_mkl_linear",
+                has_out_variant=False,
+                kernel_creator=mkldnn_ir.MKLPackedLinear.create,
+            )
+            cpu_needs_realized_inputs.append(torch.ops.mkl._mkl_linear)
+
+            @register_lowering(torch.ops.mkl._mkl_linear)
+            def mkl_packed_linear(
+                x: TensorBox,
+                packed_w: TensorBox,
+                orig_w: TensorBox,
+                b: TensorBox | None,
+                batch_size,
+                *,
+                layout=None,
+            ):
+                choices: list[ChoiceCaller] = []
+                if config.max_autotune or config.max_autotune_gemm:
+                    transposed_w = permute(orig_w, [1, 0])
+                    *_, layout, x, transposed_w = mm_args(
+                        x, transposed_w, layout=layout
+                    )
+                    if use_cpp_gemm_template(layout, x, transposed_w):
+                        CppGemmTemplate.add_choices(
+                            choices,
+                            layout,
+                            [x, packed_w, orig_w],
+                            trans_w=True,
+                            input_indices=[0, 2],
+                        )
+
+                if len(choices) == 0 or use_aten_gemm_kernels():
+                    choices.append(
+                        aten_mkl_linear.bind(
+                            (x, packed_w, orig_w), layout, B=None, batch_size=batch_size
+                        )
+                    )
+
+                assert packed_w.get_name() in V.graph.constants
+                assert orig_w.get_name() in V.graph.constants
+                # packed_w is a mkldnn tensor which we can't generate directly
+                # so we use the weights from the original tensor in autotune.
+                input_gen_fns = {
+                    1: lambda x: V.graph.constants[x.get_name()],
+                    2: lambda x: V.graph.constants[x.get_name()],
+                }
+                result: TensorBox  # annotation on separate line since tuple unpacking doesn't support inline annotation
+                result, _ = autotune_select_algorithm(
+                    "packed_linear",
+                    choices,
+                    [x, packed_w, orig_w],
+                    layout,
+                    input_gen_fns=input_gen_fns,
+                )
+                if b is not None:
+                    result = add(result, b)
+                return result
+
+        add_needs_realized_inputs(cpu_needs_realized_inputs)

outputs/audit_venv/lib/python3.11/site-packages/torch/_inductor/mock_cache.py

@@ -0,0 +1,274 @@
+# mypy: ignore-errors
+
+from __future__ import annotations
+
+import contextlib
+import dataclasses
+import sys
+import threading
+from typing import Any, TYPE_CHECKING
+from typing_extensions import override, Self
+from unittest.mock import patch
+
+from torch._inductor import config
+from torch._inductor.remote_cache import RemoteCacheBackend
+
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+    from types import TracebackType
+
+
+@dataclasses.dataclass
+class Stats:
+    num_put: int = 0
+    num_get_hit: int = 0
+    num_get_miss: int = 0
+
+    def __iadd__(self, other: Stats) -> Self:
+        self.num_put += other.num_put
+        self.num_get_hit += other.num_get_hit
+        self.num_get_miss += other.num_get_miss
+        return self
+
+    def reset(self) -> None:
+        self.num_put = 0
+        self.num_get_hit = 0
+        self.num_get_miss = 0
+
+    def __str__(self) -> str:
+        return "".join(
+            (
+                f"puts: {self.num_put}, ",
+                f"misses: {self.num_get_miss}, ",
+                f"hits: {self.num_get_hit}, ",
+            )
+        )
+
+    def __eq__(self, other: object) -> bool:
+        # Dataclass's default __eq__ checks that the types are the same so can't
+        # be used with _GlobalItemStats.
+        return (
+            isinstance(other, (Stats, _GlobalItemStats))
+            and self.num_put == other.num_put
+            and self.num_get_hit == other.num_get_hit
+            and self.num_get_miss == other.num_get_miss
+        )
+
+
+class _GlobalItemStats(Stats):
+    cache: dict[str, object]
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.cache = {}
+
+    def reset(self) -> None:
+        super().reset()
+        self.cache = {}
+
+
+# The cache states are thread-local so if we're running multiple tests at once
+# they won't cross contaminate. However - it needs to be "global" because we
+# allow code to create new cache clients which refer to the same cache (because
+# it's a remote cache).
+
+
+class _GlobalStats(threading.local):
+    def __init__(self) -> None:
+        self.autotune_local = _GlobalItemStats()
+        self.autotune_remote = _GlobalItemStats()
+        self.bundled_autotune = _GlobalItemStats()
+        self.fx_graph = _GlobalItemStats()
+        self.triton = _GlobalItemStats()
+        self.aot_autograd = _GlobalItemStats()
+        self.dynamo_pgo = _GlobalItemStats()
+
+    def reset(self) -> None:
+        self.autotune_local.reset()
+        self.autotune_remote.reset()
+        self.bundled_autotune.reset()
+        self.fx_graph.reset()
+        self.triton.reset()
+        self.aot_autograd.reset()
+        self.dynamo_pgo.reset()
+
+    def get_stat(self, name: str) -> _GlobalItemStats:
+        return getattr(self, name)
+
+    def report(self):
+        subs = (
+            ("autotune_local", self.autotune_local),
+            ("autotune_remote", self.autotune_remote),
+            ("bundled_autotune", self.bundled_autotune),
+            ("fx_graph", self.fx_graph),
+            ("triton", self.triton),
+            ("aot_autograd", self.aot_autograd),
+            ("dynamo_pgo", self.dynamo_pgo),
+        )
+
+        print("Cache Stats:", file=sys.stderr)
+        for name, sub in subs:
+            print(f"  {name}: {sub}", file=sys.stderr)
+
+        print("Cache Entries:", file=sys.stderr)
+        for name, sub in subs:
+            if sub.cache:
+                print(f"  {name}:", file=sys.stderr)
+                for k, v in sorted(sub.cache.items()):
+                    v = repr(v)
+                    if len(v) > 100:
+                        v = v[:100] + "..."
+                    print(f"    {k!r}: {v}", file=sys.stderr)
+
+
+global_stats = _GlobalStats()
+
+
+class MockBackend(RemoteCacheBackend[Any]):
+    def __init__(self, name: str) -> None:
+        self._name = name
+
+    @staticmethod
+    def with_name(name: str) -> Callable[[], MockBackend]:
+        def wrapper() -> MockBackend:
+            return MockBackend(name)
+
+        return wrapper
+
+    @override
+    def _get(self, key: str) -> Any | None:
+        stat = global_stats.get_stat(self._name)
+        if key in stat.cache:
+            stat += Stats(num_get_hit=1)
+            return stat.cache.get(key)
+        else:
+            stat += Stats(num_get_miss=1)
+            return None
+
+    @override
+    def _put(self, key: str, data: Any) -> None:
+        stat = global_stats.get_stat(self._name)
+        stat += Stats(num_put=1)
+        stat.cache[key] = data
+
+
+# List of configs for each cache
+_CACHE_CONFIG_EN = (
+    "fx_graph_cache",
+    "fx_graph_remote_cache",
+    "autotune_local_cache",
+    "autotune_remote_cache",
+    "bundled_autotune_remote_cache",
+)
+
+
+class PatchCaches(contextlib.AbstractContextManager):
+    @classmethod
+    def setUp(cls):
+        # If this test is using PatchCaches then disable all the caches by
+        # default, letting the tests turn them on explicitly. This is because
+        # tests using PatchCaches will often want to check stats explicitly.
+        cls._savedCacheState = {}
+        for name in _CACHE_CONFIG_EN:
+            if hasattr(config, name):
+                cls._savedCacheState[name] = getattr(config, name)
+            setattr(config, name, False)
+
+    @classmethod
+    def tearDown(cls):
+        # Restore cache defaults
+        for name in _CACHE_CONFIG_EN:
+            delattr(config, name)
+            if name in cls._savedCacheState:
+                setattr(config, name, cls._savedCacheState[name])
+
+    def __init__(self) -> None:
+        self._stack = contextlib.ExitStack()
+
+    def __enter__(self) -> Self:
+        global_stats.reset()
+        self._stack.__enter__()
+
+        ctx = patch(
+            "torch._inductor.runtime.autotune_cache.LocalAutotuneCache.backend_override_cls",
+            MockBackend.with_name("autotune_local"),
+        )
+        self._stack.enter_context(ctx)
+
+        ctx = patch(
+            "torch._inductor.remote_cache.RemoteAutotuneCache.backend_override_cls",
+            MockBackend.with_name("autotune_remote"),
+        )
+        self._stack.enter_context(ctx)
+
+        ctx = patch(
+            "torch._inductor.remote_cache.RemoteBundledAutotuneCache.backend_override_cls",
+            MockBackend.with_name("bundled_autotune"),
+        )
+        self._stack.enter_context(ctx)
+
+        ctx = patch(
+            "torch._inductor.remote_cache.RemoteFxGraphCache.backend_override_cls",
+            MockBackend.with_name("fx_graph"),
+        )
+        self._stack.enter_context(ctx)
+
+        ctx = patch(
+            "torch._inductor.remote_cache.RemoteAOTAutogradCache.backend_override_cls",
+            MockBackend.with_name("aot_autograd"),
+        )
+        self._stack.enter_context(ctx)
+
+        ctx = patch(
+            "torch._inductor.remote_cache.RemoteDynamoPGOCache.backend_override_cls",
+            MockBackend.with_name("dynamo_pgo"),
+        )
+        self._stack.enter_context(ctx)
+
+        if config.is_fbcode():
+            ctx = patch(
+                "torch._inductor.fb.remote_cache.FbRemoteAutotuneCache.backend_override_cls",
+                MockBackend.with_name("autotune_remote"),
+            )
+            self._stack.enter_context(ctx)
+
+            ctx = patch(
+                "torch._inductor.fb.remote_cache.FbRemoteBundledAutotuneCache.backend_override_cls",
+                MockBackend.with_name("bundled_autotune"),
+            )
+            self._stack.enter_context(ctx)
+
+            ctx = patch(
+                "torch._inductor.fb.remote_cache.FbRemoteFxGraphCache.backend_override_cls",
+                MockBackend.with_name("fx_graph"),
+            )
+            self._stack.enter_context(ctx)
+
+            ctx = patch(
+                "triton.fb.fb_memcache.FbMemcacheRemoteKernelCache.backend_override_cls",
+                MockBackend.with_name("triton"),
+            )
+            self._stack.enter_context(ctx)
+
+            ctx = patch(
+                "torch._inductor.fb.remote_cache.FbRemoteAOTAutogradCache.backend_override_cls",
+                MockBackend.with_name("aot_autograd"),
+            )
+            self._stack.enter_context(ctx)
+
+            ctx = patch(
+                "torch._inductor.fb.remote_cache.FbRemoteDynamoPGOCache.backend_override_cls",
+                MockBackend.with_name("dynamo_pgo"),
+            )
+            self._stack.enter_context(ctx)
+
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        self._stack.__exit__(exc_type, exc_value, traceback)

outputs/audit_venv/lib/python3.11/site-packages/torch/_inductor/ops_handler.py

@@ -0,0 +1,1207 @@
+# mypy: allow-untyped-defs
+from __future__ import annotations
+
+import inspect
+import itertools
+import re
+import warnings
+from io import StringIO
+from typing import Any, Generic, Literal, NamedTuple, TYPE_CHECKING, TypeVar
+from unittest.mock import patch
+
+import sympy
+
+import torch
+import torch.utils._pytree as pytree
+
+from ..utils._ordered_set import OrderedSet
+from .utils import IndentedBuffer, reduction_num_outputs, sympy_index_symbol, sympy_str
+
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
+T = TypeVar("T")
+AtomicMode = Literal[
+    "atomic_add",
+    "atomic_max",
+    "atomic_min",
+    "atomic_and",
+    "atomic_or",
+    "atomic_xor",
+    "atomic_cas",
+    "atomic_xchg",
+]
+StoreMode = AtomicMode | Literal["tma"] | None
+ReductionType = Literal[
+    "argmax",
+    "argmin",
+    "welford_reduce",
+    "welford_combine",
+    "any",
+    "max",
+    "min",
+    "prod",
+    "sum",
+    "dot",
+    "xor_sum",
+    "online_softmax_reduce",
+]
+
+
+def _arg_str(a: object) -> str:
+    if isinstance(a, sympy.Expr):
+        return sympy_str(a)
+    return str(a)
+
+
+# See OpDecompositions for superclass that desugars operations like reciprocal/square.
+class OpsHandler(Generic[T]):
+    """
+    Protocol describing the set of valid operations on ``torch._inductor.virtualized.ops``,
+    as well as the contract for op handlers.  The type T signifies the domain
+    of the abstract analysis AKA what all the functions return / take as arguments
+    anywhere compute occurs.
+
+    While these operators are typically dtype polymorphic (e.g., you can use mul
+    on both integers and floats), they do NOT do promotion and usually return the
+    same dtype as the input.  You are expected to have handled type promotion
+    during ATen decompositions.  Most operators correspond exactly to pointwise
+    operations as defined by torch, so when in doubt about semantics, check the
+    corresponding torch documentation.  These are all scalar operations (so they
+    are defined to operate on a single element at a time.)
+
+    For convenience, many operators take a src_dtype which indicates what the dtype
+    of the input argument is.  Although in principle this can be derived by an
+    analysis, providing this for ops where it is useful helps avoid having to repeatedly
+    recompute dtype in code generation.
+
+    Note that this often describes a class of static methods, for stateless
+    ops handlers.
+
+    Handlers are often defined using metaprogramming (e.g. _initialize_pointwise_overrides),
+    which means you will not get type errors for those methods.  We have tests in
+    test/inductor/test_op_completeness.py which check that all operators are implemented after
+    all the metaprogramming has run.
+    """
+
+    def constant(self, value: bool | float | int, dtype: torch.dtype) -> T:
+        """Produces a scalar constant of type dtype."""
+        raise NotImplementedError
+
+    def load_seed(self, name: str, offset: T) -> T:
+        """Computes inductor_prims.lookup_seed."""
+        raise NotImplementedError
+
+    def rand(self, seed: T, offset: T) -> T:
+        """Computes inductor_prims.random with mode="rand".  offset has dtype int32."""
+        raise NotImplementedError
+
+    def rand_eager(
+        self, seed: T, base_offset: T, threads_per_round: T, tid: T, vec: T
+    ) -> T:
+        """Computes inductor_prims.random with mode="rand_eager".  offset has dtype int32."""
+        raise NotImplementedError
+
+    def randn(self, seed: T, offset: T) -> T:
+        """Computes inductor_prims.random with mode="randn".  offset has dtype int32."""
+        raise NotImplementedError
+
+    def randint64(self, seed: T, offset: T, low: T, high: T) -> T:
+        """Computes inductor_prims.randint.  offset has dtype int32."""
+        raise NotImplementedError
+
+    def masked(self, mask: T, body: Callable[[], T], other: T) -> T:
+        """
+        Computes body, but only perform loads/stores if the boolean mask
+        evaluates to true.  For example, you would use this if you needed to
+        perform an indirect load that may not be valid on some elements;
+        without masking, invalid accesses can cause IMAs.  When mask is true,
+        the result is the result of body; otherwise it is other. Here, `other`
+        needs to be a constant.
+
+        Contrast this with ops.where, which can multiplex between two values
+        that have been unconditionally computed.
+        """
+        raise NotImplementedError
+
+    def where(self, condition: T, input: T, other: T) -> T:
+        """
+        Computes torch.where: when condition is true, return input; otherwise return other.
+        """
+        raise NotImplementedError
+
+    def index_expr(self, expr: sympy.Expr, dtype: torch.dtype) -> T:
+        """
+        Converts a sympy expression into a scalar of type dtype.  expr is typically
+        an indexing expression, thus the name; however, it can also be used in
+        non-indexing situations.
+        """
+        raise NotImplementedError
+
+    def to_dtype(
+        self,
+        x: T,
+        dtype: torch.dtype,
+        src_dtype: torch.dtype | None = None,
+        use_compute_types: bool = True,
+    ) -> T:
+        """
+        Convert x to dtype.  src_dtype can be optionally set to specify what the original
+        dtype of x was, which can improve code generation (used by torch to(dtype=dtype)).
+        """
+        raise NotImplementedError
+
+    def trunc_to_int(self, x: T, dtype: torch.dtype) -> T:
+        """
+        Convert x to dtype with truncation semantics (similar to how the int
+        constructor works in Python).  In Inductor codegen, this just decays
+        to trunc and then to_dtype, but this composite operation helps
+        roundtrips for Sympy evaluation.
+
+        dtype is taken as an explicit parameter because the desired output
+        dtype is typically the index dtype, which may vary between int32 and
+        int64 depending on if we've shown that all the indexing operations can
+        be done in int32.
+        """
+        raise NotImplementedError
+
+    def ceil_to_int(self, x: T, dtype: torch.dtype) -> T:
+        """
+        Convert x to dtype with ceiling semantics.  See also trunc_to_int.
+        """
+        raise NotImplementedError
+
+    def floor_to_int(self, x: T, dtype: torch.dtype) -> T:
+        """
+        Convert x to dtype with ceiling semantics.  See also trunc_to_int.
+        """
+        raise NotImplementedError
+
+    def round_to_int(self, x: T, dtype: torch.dtype) -> T:
+        """
+        Convert x to dtype with round-to-even semantics.  See also trunc_to_int.
+        """
+        raise NotImplementedError
+
+    def to_dtype_bitcast(self, x: T, dtype: torch.dtype, src_dtype: torch.dtype) -> T:
+        """
+        Reinterpret cast x to dtype (reinterpreting the bits in memory as another dtype.)
+        src_dtype must be the original type of x.
+        """
+        raise NotImplementedError
+
+    def identity(self, x: T) -> T:
+        """
+        Returns x as is.  This is used to trigger CSE.
+        """
+        raise NotImplementedError
+
+    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+    # These operations are only available in a "kernel" context.  Check
+    # torch._inductor.codegen.common.CSEProxy for their typical implementation
+    # in op handler (routing to their respective implementations in the kernel
+    # handler)
+    #
+    # Importantly, inside a kernel, indexing and mask variables are available
+    # in scope, which are typically used by sympy.Expr indexing.
+
+    def indirect_indexing(
+        self, x: T, size: sympy.Expr, check: bool = True, wrap_neg=True
+    ) -> sympy.Expr:
+        """
+        Convert an integral x into a sympy.Expr that can be subsequently used in
+        indexing computation.  'size' represents an upper bound on what valid
+        indexes can be; when 'check' is True, we check that the x is in bounds.
+
+        NB: This is typically mandatory to implement for any analysis, because you
+        MUST return a valid sympy.Expr of some sort (even if it's a meaningless symbol).
+        """
+        raise NotImplementedError
+
+    def load(self, name: str, index: sympy.Expr) -> T:
+        """
+        Load from the memory location 'name', offset by some indexing expression 'index'.
+        """
+        raise NotImplementedError
+
+    def store(
+        self,
+        name: str,
+        index: sympy.Expr,
+        value: T,
+        mode: StoreMode = None,
+    ) -> None:
+        """
+        Store 'value' to the memory location 'name' offset by 'expr'.  If
+        specified, 'mode' can require the store to be an atomic addition.
+        """
+        raise NotImplementedError
+
+    # TODO: Better explain how the "collective" semantics of these ops;
+    # remember that the input value is a scalar, you can't reduce on it in the
+    # traditional sense!
+    def reduction(
+        self,
+        dtype: torch.dtype,
+        src_dtype: torch.dtype,
+        reduction_type: ReductionType,
+        value: T,
+    ) -> T | tuple[T, ...]:
+        """
+        Perform a 'reduction_type' reduction on 'value' of dtype 'src_dtype',
+        using 'dtype' as the accumulation dtype for the reduction.  The result
+        is an intermediate computation which should be stored to the final
+        location using 'ops.store_reduction'.
+
+        Valid reduction types are .  For Welford reduction types, this
+        function returns multiple outputs; consult reduction_num_outputs to
+        determine the amount in metaprogramming applications.
+        """
+        raise NotImplementedError
+
+    # TODO: in practice, this seems to actually return None, but not returning
+    # a T makes common __getattr__ idioms not type correctly.  Figure out if
+    # this should be returning something.
+    def store_reduction(self, name: str, index: sympy.Expr, value: T) -> None:
+        """
+        Store the fully accumulated result of 'reduction' to the memory
+        location 'name' offset by 'expr'.
+        """
+        raise NotImplementedError
+
+    def scan(
+        self,
+        dtypes: tuple[torch.dtype, ...],
+        combine_fn: Callable[[tuple[T, ...], tuple[T, ...]], tuple[T, ...]],
+        values: tuple[T, ...],
+    ) -> tuple[T, ...]:
+        """
+        Perform an associative scan on 'value'.
+        """
+        # TODO: Improve the description with some pseudocode
+        raise NotImplementedError
+
+    def sort(
+        self,
+        dtypes: tuple[torch.dtype, ...],
+        values: tuple[T, ...],
+        stable: bool,
+        descending: bool,
+    ) -> tuple[T, ...]:
+        """
+        Sort values along the reduction dimension.
+        """
+        raise NotImplementedError
+
+    def bucketize(
+        self,
+        values: T,
+        boundaries: tuple[str, sympy.Expr, sympy.Expr, sympy.Expr],
+        boundary_indices: T,
+        indexing_dtype: torch.dtype,
+        right: bool,
+        sorter: tuple[str, sympy.Expr] | None = None,
+        sorter_indices: T | None = None,
+    ) -> T:
+        # See [Note: Inductor bucketize op]
+        raise NotImplementedError
+
+    def partial_accumulate(
+        self,
+        name: str,
+        reduction_type: ReductionType,
+        value: T,
+        extra_meta: dict[str, Any],
+    ) -> None:
+        raise NotImplementedError
+
+    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+    # The following ops have semantics that correspond exactly to the torch
+    # operation with the same corresponding name.
+
+    def abs(self, x0: T) -> T:
+        raise NotImplementedError
+
+    def exp(self, x0: T) -> T:
+        raise NotImplementedError
+
+    def exp2(self, x0: T) -> T:
+        raise NotImplementedError
+
+    def expm1(self, x0: T) -> T:
+        raise NotImplementedError
+
+    def sqrt(self, x0: T) -> T:
+        raise NotImplementedError
+
+    def relu(self, x0: T) -> T:
+        raise NotImplementedError
+
+    def minimum(self, x0: T, x1: T) -> T:
+        raise NotImplementedError
+
+    def maximum(self, x0: T, x1: T) -> T:
+        raise NotImplementedError
+
+    def cos(self, x0: T) -> T:
+        raise NotImplementedError
+
+    def sin(self, x0: T) -> T:
+        raise NotImplementedError
+
+    def lgamma(self, x0: T) -> T:
+        raise NotImplementedError
+
+    def erf(self, x0: T) -> T:
+        raise NotImplementedError
+
+    def cosh(self, x0: T) -> T:
+        raise NotImplementedError
+
+    def sinh(self, x0: T) -> T:
+        raise NotImplementedError
+
+    def acos(self, x0: T) -> T:
+        raise NotImplementedError
+
+    def acosh(self, x0: T) -> T:
+        raise NotImplementedError
+
+    def asin(self, x0: T) -> T:
+        raise NotImplementedError
+
+    def asinh(self, x0: T) -> T:
+        raise NotImplementedError
+
+    def atan2(self, x0: T, x1: T) -> T:
+        raise NotImplementedError
+
+    def atan(self, x0: T) -> T:
+        raise NotImplementedError
+
+    def atanh(self, x0: T) -> T:
+        raise NotImplementedError
+
+    def copysign(self, x0: T, x1: T) -> T:
+        raise NotImplementedError
+
+    def erfc(self, x0: T) -> T:
+        raise NotImplementedError
+
+    def erfinv(self, x0: T) -> T:
+        raise NotImplementedError
+
+    def frexp(self, x0: T):
+        raise NotImplementedError
+
+    def hypot(self, x0: T, x1: T) -> T:
+        raise NotImplementedError
+
+    def log10(self, x0: T) -> T:
+        raise NotImplementedError
+
+    def log2(self, x0: T) -> T:
+        raise NotImplementedError
+
+    def ldexp(self, x0: T, n: T) -> T:
+        raise NotImplementedError
+
+    def nextafter(self, x0: T, x1: T) -> T:
+        raise NotImplementedError
+
+    def logical_and(self, x0: T, x1: T) -> T:
+        raise NotImplementedError
+
+    def logical_not(self, x0: T) -> T:
+        raise NotImplementedError
+
+    def logical_or(self, x0: T, x1: T) -> T:
+        raise NotImplementedError
+
+    def logical_xor(self, x0: T, x1: T) -> T:
+        raise NotImplementedError
+
+    def bitwise_and(self, x0: T, x1: T) -> T:
+        raise NotImplementedError
+
+    def bitwise_not(self, x0: T) -> T:
+        raise NotImplementedError
+
+    def bitwise_or(self, x0: T, x1: T) -> T:
+        raise NotImplementedError
+
+    def bitwise_xor(self, x0: T, x1: T) -> T:
+        raise NotImplementedError
+
+    def bitwise_left_shift(self, x0: T, x1: T) -> T:
+        raise NotImplementedError
+
+    def bitwise_right_shift(self, x0: T, x1: T) -> T:
+        raise NotImplementedError
+
+    def rsqrt(self, x0: T) -> T:
+        raise NotImplementedError
+
+    def log1p(self, x0: T) -> T:
+        raise NotImplementedError
+
+    def tan(self, x0: T) -> T:
+        raise NotImplementedError
+
+    def tanh(self, x0: T) -> T:
+        raise NotImplementedError
+
+    def sigmoid(self, x0: T) -> T:
+        raise NotImplementedError
+
+    def signbit(self, x0: T) -> T:
+        raise NotImplementedError
+
+    def fmod(self, x0: T, x1: T) -> T:
+        raise NotImplementedError
+
+    def log(self, x0: T) -> T:
+        raise NotImplementedError
+
+    def isinf(self, x0: T) -> T:
+        raise NotImplementedError
+
+    def isnan(self, x0: T) -> T:
+        raise NotImplementedError
+
+    # NB: this returns a float, like the torch operation
+    # This rounds half to even to break ties
+    def round(self, x0: T) -> T:
+        raise NotImplementedError
+
+    # NB: this returns a float, like the torch operation
+    def floor(self, x0: T) -> T:
+        raise NotImplementedError
+
+    def sign(self, x0: T) -> T:
+        raise NotImplementedError
+
+    # NB: this returns a float, like the torch operation
+    def trunc(self, x0: T) -> T:
+        raise NotImplementedError
+
+    # NB: this returns a float, like the torch operation
+    def ceil(self, x0: T) -> T:
+        raise NotImplementedError
+
+    def neg(self, x0: T) -> T:
+        raise NotImplementedError
+
+    def reciprocal(self, x0: T) -> T:
+        raise NotImplementedError
+
+    def eq(self, x0: T, x1: T) -> T:
+        raise NotImplementedError
+
+    def ne(self, x0: T, x1: T) -> T:
+        raise NotImplementedError
+
+    def lt(self, x0: T, x1: T) -> T:
+        raise NotImplementedError
+
+    def gt(self, x0: T, x1: T) -> T:
+        raise NotImplementedError
+
+    def le(self, x0: T, x1: T) -> T:
+        raise NotImplementedError
+
+    def ge(self, x0: T, x1: T) -> T:
+        raise NotImplementedError
+
+    def add(self, x0: T, x1: T) -> T:
+        raise NotImplementedError
+
+    def sub(self, x0: T, x1: T) -> T:
+        raise NotImplementedError
+
+    def mul(self, x0: T, x1: T) -> T:
+        raise NotImplementedError
+
+    # NB: this returns a float, like the torch operation
+    def pow(self, x0: T, x1: T) -> T:
+        raise NotImplementedError
+
+    def and_(self, x0: T, x1: T) -> T:
+        raise NotImplementedError
+
+    def or_(self, x0: T, x1: T) -> T:
+        raise NotImplementedError
+
+    def xor(self, x0: T, x1: T) -> T:
+        raise NotImplementedError
+
+    # These are metaprogrammed by MockHandler._init_cls
+    def lshift(self, x0: T, x1: T) -> T:
+        raise NotImplementedError
+
+    def rshift(self, x0: T, x1: T) -> T:
+        raise NotImplementedError
+
+    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+    # These are "special" operators.  These only exist if the target
+    # language actually supports the operator.  Keep this in sync with
+    # pointwise_overrides_data.
+
+    def airy_ai(self, x: T) -> T:
+        raise NotImplementedError
+
+    def bessel_j0(self, x: T) -> T:
+        raise NotImplementedError
+
+    def bessel_j1(self, x: T) -> T:
+        raise NotImplementedError
+
+    def bessel_y0(self, x: T) -> T:
+        raise NotImplementedError
+
+    def bessel_y1(self, x: T) -> T:
+        raise NotImplementedError
+
+    def digamma(self, x: T) -> T:
+        raise NotImplementedError
+
+    def erfcx(self, x: T) -> T:
+        raise NotImplementedError
+
+    def fma(self, x: T, y: T, z: T) -> T:
+        raise NotImplementedError
+
+    def mul_rn(self, x: T, y: T) -> T:
+        """Multiplication with round-to-nearest, preventing fusion with subsequent ops."""
+        raise NotImplementedError
+
+    def igamma(self, x: T, y: T) -> T:
+        raise NotImplementedError
+
+    def igammac(self, x: T, y: T) -> T:
+        raise NotImplementedError
+
+    def gammainc(self, x: T, y: T) -> T:
+        raise NotImplementedError
+
+    def gammaincc(self, x: T, y: T) -> T:
+        raise NotImplementedError
+
+    def i0(self, x: T) -> T:
+        raise NotImplementedError
+
+    def i0e(self, x: T) -> T:
+        raise NotImplementedError
+
+    def i1(self, x: T) -> T:
+        raise NotImplementedError
+
+    def i1e(self, x: T) -> T:
+        raise NotImplementedError
+
+    def log_ndtr(self, x: T) -> T:
+        raise NotImplementedError
+
+    def modified_bessel_i0(self, x: T) -> T:
+        raise NotImplementedError
+
+    def modified_bessel_i1(self, x: T) -> T:
+        raise NotImplementedError
+
+    def modified_bessel_k0(self, x: T) -> T:
+        raise NotImplementedError
+
+    def modified_bessel_k1(self, x: T) -> T:
+        raise NotImplementedError
+
+    def ndtr(self, x: T) -> T:
+        raise NotImplementedError
+
+    def ndtri(self, x: T) -> T:
+        raise NotImplementedError
+
+    def polygamma(self, x: T, y: T) -> T:
+        raise NotImplementedError
+
+    def scaled_modified_bessel_k0(self, x: T) -> T:
+        raise NotImplementedError
+
+    def scaled_modified_bessel_k1(self, x: T) -> T:
+        raise NotImplementedError
+
+    def spherical_bessel_j0(self, x: T) -> T:
+        raise NotImplementedError
+
+    def zeta(self, x: T, y: T) -> T:
+        raise NotImplementedError
+
+    def chebyshev_polynomial_t(self, x: T, y: T) -> T:
+        raise NotImplementedError
+
+    def chebyshev_polynomial_u(self, x: T, y: T) -> T:
+        raise NotImplementedError
+
+    def chebyshev_polynomial_v(self, x: T, y: T) -> T:
+        raise NotImplementedError
+
+    def chebyshev_polynomial_w(self, x: T, y: T) -> T:
+        raise NotImplementedError
+
+    def legendre_polynomial_p(self, x: T, y: T) -> T:
+        raise NotImplementedError
+
+    def shifted_chebyshev_polynomial_t(self, x: T, y: T) -> T:
+        raise NotImplementedError
+
+    def shifted_chebyshev_polynomial_u(self, x: T, y: T) -> T:
+        raise NotImplementedError
+
+    def shifted_chebyshev_polynomial_v(self, x: T, y: T) -> T:
+        raise NotImplementedError
+
+    def shifted_chebyshev_polynomial_w(self, x: T, y: T) -> T:
+        raise NotImplementedError
+
+    def hermite_polynomial_h(self, x: T, y: T) -> T:
+        raise NotImplementedError
+
+    def hermite_polynomial_he(self, x: T, y: T) -> T:
+        raise NotImplementedError
+
+    def laguerre_polynomial_l(self, x: T, y: T) -> T:
+        raise NotImplementedError
+
+    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+    # These operators are a bit special, because they are conventionally
+    # natively supported in both Python and C, but the semantics differ so
+    # care must be taken
+
+    def truncdiv(self, x0: T, x1: T) -> T:
+        """C-style trunc division between integers only.  Computes the true
+        division of two numbers and rounds the result to zero.
+        """
+        raise NotImplementedError
+
+    def floordiv(self, x0: T, x1: T) -> T:
+        """Python-style floor division between integers only.  Computes the
+        true division of two numbers and floors the result.  If you want
+        floor division for floats, do regular truediv and floor the result.
+        """
+        raise NotImplementedError
+
+    def truediv(self, x0: T, x1: T) -> T:
+        """True division between floats.  Integer inputs are NOT valid.  To
+        do Python-style (int, int) -> float division, use int_truediv"""
+        raise NotImplementedError
+
+    def div_rn(self, x0: T, x1: T) -> T:
+        """Division with round-to-nearest rounding mode.  Used for matching
+        eager CUDA semantics where division uses IEEE round-to-nearest."""
+        raise NotImplementedError
+
+    def int_truediv(self, x0: T, x1: T) -> T:
+        """True division between integers.  This is NOT the same as promoting
+        to float and doing integer division, there is a bespoke algorithm for
+        doing the division in higher precision than the above.
+        """
+        raise NotImplementedError
+
+    def mod(self, x0: T, x1: T) -> T:
+        """C-style modulus, take sign from LHS (x0)."""
+        raise NotImplementedError
+
+    def remainder(self, x0: T, x1: T) -> T:
+        """Python-style modulus, take sign from RHS (x1)."""
+        raise NotImplementedError
+
+    def square(self, x0: T) -> T:
+        raise NotImplementedError
+
+    def check_bounds(
+        self, expr: sympy.Expr, size: sympy.Expr, lower: bool, upper: bool
+    ) -> None:
+        raise NotImplementedError
+
+    # halide-only
+    def halide_clamp(self, value: T, size: sympy.Expr, check: bool) -> T:
+        raise NotImplementedError
+
+    # triton-only
+    def dot(self, x: T, y: T) -> T:
+        raise NotImplementedError
+
+    # triton-only
+    def inline_asm_elementwise(
+        self,
+        *inputs: T,
+        asm: str,
+        constraints: str | None = None,
+        dtype: torch.dtype = torch.float32,
+        is_pure: bool = True,
+        pack: int = 1,
+        input_dtypes: tuple[torch.dtype, ...] | None = None,
+    ) -> T:
+        raise NotImplementedError
+
+    def output(self, *args: T) -> None:
+        """This is a fake op used in analysis but not codegen"""
+        raise NotImplementedError
+
+    def placeholder(self, index: int) -> T:
+        """This is a fake op used in analysis but not codegen"""
+        raise NotImplementedError
+
+    def device_assert_async(self, cond: T, msg: str) -> T:
+        raise NotImplementedError
+
+
+_ignore_op_re = re.compile(r"_.*|paren").fullmatch
+
+
+def list_ops(cls: type[Any]):
+    return OrderedSet([x for x in dir(cls) if not _ignore_op_re(x)])
+
+
+OP_NAMES = list_ops(OpsHandler)
+
+
+class DefaultHandler(OpsHandler[Any]):
+    def _default(self, name: str, args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any:
+        """
+        Default implementation for all ops.  Override in a subclass to
+        provide generic op behavior.
+
+        Args:
+            name: name of the op, see OpHandler.{name}
+            args: positional args passed to the op
+            kwargs: keyword args passed to the op
+
+        Returns:
+            return value of the op
+
+        """
+        raise NotImplementedError
+
+    def __getattr__(self, name: str) -> Any:
+        def fallback(*args: Any, **kwargs: Any) -> Any:
+            return self._default(name, args, kwargs)
+
+        # would like to remove this function entirely, but it's used in MTIA backend
+        warnings.warn(f"undefined OpHandler.{name}, please add missing op schema")
+        return fallback
+
+    @staticmethod
+    def _call_default(target: str):
+        def call_default(self, *args, **kwargs):
+            return self._default(target, args, kwargs)
+
+        call_default.__name__ = target
+        return call_default
+
+    @classmethod
+    def _init_cls(cls):
+        """
+        Here we codegen many functions of the form:
+
+            def add(self, a, b):
+                return self._default('add', (a, b), {})
+
+        and install them in cls.  This is the same as _call_default above,
+        but is about 1.2x faster since CPython varargs parsing is slow.
+        """
+        code = StringIO()
+        for target in OP_NAMES:
+            sig = inspect.signature(getattr(OpsHandler, target))
+            if all(
+                p.kind == inspect.Parameter.POSITIONAL_OR_KEYWORD
+                and p.default is inspect.Parameter.empty
+                for p in sig.parameters.values()
+            ):
+                self_arg, *args = sig.parameters.keys()
+                assert self_arg == "self"
+                code.write(
+                    f"""
+                    def {target}(self, {", ".join(args)}):
+                        return self._default({target!r}, ({", ".join(args)}, ), {{}})
+                    """.strip()
+                )
+                code.write("\n\n")
+            else:
+                # slower fallback for ops with default or variadic arguments
+                setattr(cls, target, cls._call_default(target))
+
+        ctx: dict[str, Any] = {}
+        exec(code.getvalue(), ctx)
+        for target, impl in ctx.items():
+            if target in OP_NAMES:
+                setattr(cls, target, impl)
+
+
+DefaultHandler._init_cls()
+
+
+class NoopHandler(DefaultHandler):
+    name = "NoopHandler"
+
+    def _default(self, name: str, args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any:
+        return None
+
+    @staticmethod
+    def masked(mask, body, other) -> None:
+        return None
+
+    @staticmethod
+    # pyrefly: ignore [bad-override]
+    def frexp(x) -> tuple[None, None]:
+        return (None, None)
+
+    @staticmethod
+    def scan(dtypes, combine_fn, values) -> tuple[None, ...]:
+        return (None,) * len(values)
+
+    @staticmethod
+    def sort(dtypes, values, stable, descending) -> tuple[None, ...]:
+        return (None,) * len(values)
+
+    @staticmethod
+    # pyrefly: ignore [bad-override]
+    def indirect_indexing(index_var, size, check=True, wrap_neg=True) -> sympy.Symbol:
+        return sympy.S.Zero
+
+
+class BasicMathOpsMixin:
+    @staticmethod
+    def add(a, b):
+        return f"{a} + {b}"
+
+    @staticmethod
+    def sub(a, b):
+        return f"{a} - {b}"
+
+    @staticmethod
+    def mul(a, b):
+        return f"{a} * {b}"
+
+    @staticmethod
+    def floordiv(a, b):
+        return f"{a} // {b}"
+
+    @staticmethod
+    def truediv(a, b):
+        return f"{a} / {b}"
+
+    @staticmethod
+    def mod(a, b):
+        # careful, depending on target semantics varies
+        return f"{a} % {b}"
+
+    @staticmethod
+    def pow(a, b):
+        return f"{a} ** {b}"
+
+    @staticmethod
+    def lshift(a, b):
+        return f"{a} << {b}"
+
+    @staticmethod
+    def rshift(a, b):
+        return f"{a} >> {b}"
+
+    @staticmethod
+    def and_(a, b):
+        return f"{a} & {b}"
+
+    @staticmethod
+    def or_(a, b):
+        return f"{a} | {b}"
+
+    @staticmethod
+    def xor(a, b):
+        return f"{a} ^ {b}"
+
+    @staticmethod
+    def eq(a, b):
+        return f"{a} == {b}"
+
+    @staticmethod
+    def ne(a, b):
+        return f"{a} != {b}"
+
+    @staticmethod
+    def lt(a, b):
+        return f"{a} < {b}"
+
+    @staticmethod
+    def gt(a, b):
+        return f"{a} > {b}"
+
+    @staticmethod
+    def le(a, b):
+        return f"{a} <= {b}"
+
+    @staticmethod
+    def ge(a, b):
+        return f"{a} >= {b}"
+
+    @staticmethod
+    def neg(a):
+        return f"-{a}"
+
+
+class MockHandler(BasicMathOpsMixin, DefaultHandler):
+    name = "MockHandler"
+
+    def _default(self, name: str, args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any:
+        fargs = [*map(_arg_str, args)]
+        for k, v in kwargs.items():
+            fargs.append(f"{k}={_arg_str(v)}")
+        return f"ops.{name}({', '.join(fargs)})"
+
+    @staticmethod
+    def masked(mask, body, other) -> str:
+        return f"ops.masked({mask}, {body()}, {other})"
+
+    @staticmethod
+    # pyrefly: ignore [bad-override]
+    def frexp(x):
+        return (f"ops.frexp({x})[0]", f"ops.frexp({x})[1]")
+
+    @staticmethod
+    def scan(dtypes, combine_fn, values):
+        return tuple(
+            f"ops.scan({dtypes}, {combine_fn}, {values})[{i}]"
+            for i in range(len(values))
+        )
+
+    @staticmethod
+    def sort(dtypes, values, stable, descending):
+        return tuple(
+            f"ops.sort({dtypes}, {values}, stable={stable}, descending={descending})[{i}]"
+            for i in range(len(values))
+        )
+
+    @staticmethod
+    # pyrefly: ignore [bad-override]
+    def indirect_indexing(index_var, size, check=True, wrap_neg=True) -> sympy.Symbol:
+        return sympy_index_symbol(str(index_var))
+
+
+class KernelFormatterHandler(DefaultHandler):
+    def __init__(self, parent_handler: OpsHandler[Any]):
+        self.parent_handler = parent_handler
+        self._output = IndentedBuffer(1)
+        self.var_counter = itertools.count()
+
+    @staticmethod
+    def ir_to_string(ir_fn, index, rindex=None) -> str:
+        from .ir import FlexibleLayout
+        from .virtualized import V
+
+        args = [index, rindex] if rindex is not None else [index]
+        names = ["index", "rindex"] if rindex is not None else ["index"]
+        formatter = KernelFormatterHandler(MockHandler())
+
+        with formatter._output.indent(-1):
+            formatter._output.writeline(f"def inner_fn({', '.join(names)}):")
+        for name, arg in zip(names, args):
+            if arg:
+                lhs = ", ".join(
+                    [
+                        str("_" if isinstance(v, (int, sympy.Integer)) else v)
+                        for v in arg
+                    ]
+                )
+                formatter._output.writeline(f"{lhs} = {name}")
+
+        with (
+            V.set_ops_handler(formatter),
+            patch.object(FlexibleLayout, "allow_indexing", True),
+        ):
+            result = ir_fn(*args)
+            return formatter.getvalue(result)
+
+    def indirect_indexing(self, *args, **kwargs) -> sympy.Symbol:
+        return self.parent_handler.indirect_indexing(*args, **kwargs)
+
+    def _write(self, line):
+        # replace line with a new variable name
+        varname = f"tmp{next(self.var_counter)}"
+        self._output.writeline(f"{varname} = {line}")
+        return varname
+
+    def _default(self, name: str, args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any:
+        return pytree.tree_map(
+            self._write, getattr(self.parent_handler, name)(*args, **kwargs)
+        )
+
+    def reduction(
+        self,
+        dtype: torch.dtype,
+        src_dtype: torch.dtype,
+        reduction_type: ReductionType,
+        value: str | tuple[str, ...],
+    ) -> str | tuple[str, ...]:
+        line = self.parent_handler.reduction(dtype, src_dtype, reduction_type, value)
+        num_values = reduction_num_outputs(reduction_type)
+        varnames = [f"tmp{next(self.var_counter)}" for _ in range(num_values)]
+        self._output.writeline(f"{','.join(varnames)} = {line}")
+        return tuple(varnames) if num_values > 1 else varnames[0]
+
+    def getvalue(self, result):
+        self._output.writeline(f"return {result}")
+        return self._output.getvalue()
+
+
+class WrapperHandler(DefaultHandler):
+    def __init__(self, inner: OpsHandler[Any]):
+        self._inner = inner
+
+    def _default(self, name: str, args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any:
+        return getattr(self._inner, name)(*args, **kwargs)
+
+
+class AddParenHandler(WrapperHandler):
+    def _default(self, name: str, args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any:
+        val = getattr(self._inner, name)(*args, **kwargs)
+        if not val or isinstance(val, (sympy.Expr, tuple, list)):
+            return val
+        return f"({val})"
+
+
+class OpCountResult(NamedTuple):
+    num_ops: int
+    used_ops: OrderedSet[str]
+    read_buffers: list[str]
+    nontrivial_read_count: int
+
+
+class OpCounterCSE(DefaultHandler):
+    """Shim to count how many ops are used"""
+
+    def __init__(self, inner: OpsHandler[Any]):
+        super().__init__()
+        self.parent_handler = inner
+        self.op_count = 0
+        self.var_names: dict[str, str] = {}
+        self._used_ops: OrderedSet[str] = OrderedSet()
+        self._read_names: list[str] = []
+        self._nontrivial_read_count = 0
+
+    def _default(self, name: str, args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any:
+        self._used_ops.add(name)
+        return pytree.tree_map(
+            self._update_count, getattr(self.parent_handler, name)(*args, **kwargs)
+        )
+
+    def _update_count(self, val):
+        varname = self.var_names.get(val)
+        if not varname:
+            varname = f"tmp{self.op_count}"
+            self.op_count += 1
+            self.var_names[val] = varname
+        return varname
+
+    def indirect_indexing(self, *args, **kwargs):
+        self._used_ops.add("indirect_indexing")
+        return self.parent_handler.indirect_indexing(*args, **kwargs)
+
+    def load(self, name: str, index: sympy.Expr) -> str:
+        val = self.parent_handler.load(name, index)
+        if val not in self.var_names:
+            self._used_ops.add("load")
+            self._read_names.append(name)
+            if not isinstance(index, (sympy.Integer, int)):
+                self._nontrivial_read_count += 1
+        return self._update_count(val)
+
+    def load_seed(self, name: str, offset: T):
+        val = self.parent_handler.load_seed(name, offset)
+        if val not in self.var_names:
+            self._used_ops.add("load_seed")
+            self._read_names.append(name)
+        return self._update_count(val)
+
+    def bucketize(
+        self,
+        values: T,
+        boundaries: tuple[str, sympy.Expr, sympy.Expr, sympy.Expr],
+        boundary_indices: T,
+        indexing_dtype: torch.dtype,
+        right: bool,
+        sorter: tuple[str, sympy.Expr] | None = None,
+        sorter_indices: T | None = None,
+    ) -> T:
+        """
+        See [Note: Inductor bucketize op]
+        """
+        val = self.parent_handler.bucketize(
+            values,
+            boundaries,
+            boundary_indices,
+            indexing_dtype,
+            right,
+            sorter,
+            sorter_indices,
+        )
+        if val not in self.var_names:
+            self._used_ops.add("bucketize")
+            self._read_names.append(boundaries[0])
+            if sorter is not None:
+                self._read_names.append(sorter[0])
+        return self._update_count(val)
+
+    def getvalue(self):
+        return OpCountResult(
+            self.op_count, self._used_ops, self._read_names, self._nontrivial_read_count
+        )
+
+
+class ExtractConstantsHandler(NoopHandler):
+    def __init__(self, device: torch.device | None):
+        self.device = device
+
+    def constant(self, value: Any, dtype: torch.dtype) -> torch._inductor.ir.Constant:
+        from torch._inductor import ir
+
+        return ir.Constant(
+            value=value, dtype=dtype, device=self.device or torch.get_default_device()
+        )
+
+
+class SimpleCSEHandler(WrapperHandler):
+    """Wraps the underlying handler with a CSE pass
+
+    NOTE: Compared to codegen level CSE this is simplified as it
+    doesn't support stores which require load cache invalidation.
+    """
+
+    def __init__(self, inner: Any):
+        super().__init__(inner)
+        self.cse_cache: dict[str, Any | tuple[Any, ...]] = {}
+        self.mock = MockHandler()
+
+    def indirect_indexing(self, *args, **kwargs) -> sympy.Expr:
+        return super().indirect_indexing(*args, **kwargs)  # type: ignore[misc]
+
+    def store(self, *args, **kwargs) -> None:
+        raise NotImplementedError("store not implemented")
+
+    def store_reduction(self, *args, **kwargs) -> None:
+        raise NotImplementedError("store not implemented")
+
+    def _default(self, name: str, args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any:
+        key = getattr(self.mock, name)(*args, **kwargs)
+        val = self.cse_cache.get(key)
+        if val is not None:
+            return val
+
+        val = getattr(self._inner, name)(*args, **kwargs)
+        self.cse_cache[key] = val
+        return val
+
+    def device_assert_async(self, *args, **kwargs) -> None:
+        raise NotImplementedError(
+            f"{type(self).__name__}: device_assert_async should be handled by CSEProxy"
+        )

outputs/audit_venv/lib/python3.11/site-packages/torch/_inductor/optimize_indexing.py

@@ -0,0 +1,126 @@
+import math
+from typing import Any
+
+import sympy
+
+import torch
+from torch.utils._sympy.value_ranges import ValueRanges
+
+from .loop_body import LoopBody
+from .utils import dominated_nodes
+
+
+def val_expressable_in_32_bits(val: Any) -> bool:
+    if getattr(val, "is_Boolean", False):
+        return True
+
+    if isinstance(val, sympy.Expr):
+        assert val.is_number
+        if val.is_Integer or val.is_Boolean:
+            val = int(val)
+        else:
+            val = float(val)
+
+    # bound within mantissa
+    if isinstance(val, float):
+        return val <= (2**24) and val >= -(2**24)
+
+    if isinstance(val, int):
+        iinfo = torch.iinfo(torch.int32)
+        return val <= iinfo.max and val >= iinfo.min
+
+    raise TypeError(f"Unexpected value {val}")
+
+
+def range_expressable_in_32_bits(range: ValueRanges[sympy.Expr]) -> bool:
+    return val_expressable_in_32_bits(range.lower) and val_expressable_in_32_bits(
+        range.upper
+    )
+
+
+def try_to_reduce_precision(
+    node: Any,
+    bounds: dict[Any, Any],
+    indirect_vars: list[Any],
+    indices: dict[Any, sympy.Expr],
+    replacement_vals: dict[Any, ValueRanges[sympy.Expr]],
+) -> None:
+    # if a downstream use of a node explicitly converts to int32, or float16/float32/float64,
+    # then it's precision is set for that chain of uses, and we don't need to consider those
+    # dominated values
+    def skip_filter(node: Any) -> bool:
+        return node.target == "to_dtype" and node.args[2] in (
+            torch.int32,
+            torch.float32,
+            torch.float64,
+        )
+
+    # TODO - there are dominated uses whose dtype does not depend on whether
+    # we reduce the precision here, e.g. add(int64, int64) one of the args can be reduced to
+    # int32 without changing the output precision of the node. this case hasn't shown up
+    for dominated in dominated_nodes([node], skip_filter):
+        if dominated.target in ["store", "output"]:
+            continue
+
+        if isinstance(dominated.target, str) and "set_indirect" in dominated.target:
+            idx = int(dominated.target[len("set_indirect") :])
+            indirect_var = indirect_vars[idx]
+
+            # We check that we can compute all the indices it's involved in with int32
+            for index, expr in indices.items():
+                if indirect_var in expr.free_symbols:
+                    index_val = replacement_vals[index]
+
+                    if math.isinf(index_val.lower) or math.isinf(index_val.upper):
+                        return
+
+                    # all indices are integers, so make sure that we
+                    # use the bounds of integers instead of floats.
+                    # TODO - not sure if we should be doing int/float casts while tracing,
+                    # might interfere with sympy.
+
+                    index_val_int = ValueRanges[sympy.Expr](
+                        int(index_val.lower), int(index_val.upper)
+                    )
+                    if not range_expressable_in_32_bits(index_val_int):
+                        return
+
+        if not range_expressable_in_32_bits(bounds[dominated]):
+            return
+
+    args = list(node.args)
+    args[2] = torch.int32
+    node.args = tuple(args)
+
+
+def indexing_dtype_strength_reduction(loop_body: LoopBody) -> None:
+    """
+    Performs Value Range Analysis on LoopBody's fx graph to reduce precision of
+    intermediaries from int64 to int32
+    """
+    bv = loop_body.bounds()
+
+    int64_dtype_nodes = [
+        node
+        for node in loop_body.get_nodes()
+        if (
+            node.target == "to_dtype"
+            and node.args[2] == torch.int64
+            and node not in bv.unbounded_vars
+        )
+    ]
+    if not int64_dtype_nodes:
+        return
+
+    bounds = bv.get_bounds()
+
+    # TODO - if dominated node of one to_dtype is not expressible in int32,
+    # we should short circuit another to_dtype node if that node also dominates
+    for node in int64_dtype_nodes:
+        try_to_reduce_precision(
+            node,
+            bounds,
+            loop_body.indirect_vars,
+            loop_body.indexing_exprs,
+            bv.replacement_vals,
+        )

outputs/audit_venv/lib/python3.11/site-packages/torch/_inductor/output_code.py

@@ -0,0 +1,1202 @@
+"""
+This provides an abstract class which parametrizes over an "output code" concept
+for Inductor.  Intuitively, this represents the compiled callable which Inductor
+produces which you can call to get optimized code.  However, this callable
+has some other capabilities:
+
+- It is serializable, so you can save/load this product from disk without
+  having to do compilation again.
+
+- (When using remote cache) it is addressable, so you can save just a key
+  which you can use to load this product from remote cache later.
+
+This class is abstract because we have several different implementations of
+serialized format:
+
+- Python wrapper (the default)
+
+- AOTInductor (this produces ABI stable binaries which work across PyTorch
+  versions)
+
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import logging
+import os
+from functools import partial
+from typing import Any, TYPE_CHECKING, TypeAlias
+
+import torch
+from torch._dynamo.utils import counters, get_runtime_metrics_context
+from torch._higher_order_ops.wrap import inductor_compiled_code
+from torch._inductor.cudagraph_utils import (
+    BoxedDeviceIndex,
+    CudagraphCachedInfo,
+    CudagraphMetadata,
+    get_partition_cudagraph_metadata,
+    get_placeholder_info,
+    log_cudagraph_skip_and_bump_counter,
+)
+from torch._inductor.freezing_utils import has_frozen_params, is_frozen_param
+from torch._inductor.utils import (
+    _unstable_customized_partition_wrapper,
+    align_inputs_from_check_idxs,
+    BoxedBool,
+    CUDAGraphWrapperMetadata,
+    GraphPartitionMap,
+    InputType,
+    output_node,
+    set_tracing_context_output_strides,
+)
+from torch._opaque_base import OpaqueBase
+from torch.fx._graph_pickler import _node_metadata_key_filter_safe, _ops_filter_safe
+from torch.utils._ordered_set import OrderedSet
+from torch.utils._python_dispatch import is_in_torch_dispatch_mode
+
+from . import config
+from .runtime.autotune_cache import AutotuneCacheBundler
+
+
+if TYPE_CHECKING:
+    from collections import Counter
+    from collections.abc import Callable, Sequence
+
+    from torch._inductor import metrics
+    from torch._inductor.graph import GraphLowering
+    from torch._library.fake_class_registry import FakeScriptObject
+    from torch.export.pt2_archive._package_weights import Weights
+
+    from .compile_fx import _CompileFxKwargs
+    from .triton_bundler import TritonBundle
+
+log = logging.getLogger(__name__)
+
+
+@dataclasses.dataclass
+class OutputCode:
+    # TODO: Remove underscores here
+
+    # None if the output is not remote cacheable
+    _fx_graph_cache_key: str | None = dataclasses.field(default=None, init=False)
+    _fx_graph_cache_debug_lines: list[str] | None = dataclasses.field(
+        default=None, init=False
+    )
+
+    # How long it took to compile this OutputCode, end to end
+    _time_taken_ns: int | None = dataclasses.field(default=None, init=False)
+
+    def __call__(self, inputs: Sequence[Any]) -> Any:
+        raise NotImplementedError(type(self))
+
+    def prepare_for_serialization(self) -> None:
+        raise NotImplementedError(type(self))
+
+    def post_compile(
+        self,
+        example_inputs: Sequence[InputType],
+        constants: CompiledFxGraphConstants,
+        graph_kwargs: _CompileFxKwargs,
+    ) -> None:
+        raise NotImplementedError(type(self))
+
+    # TODO: Get rid of this
+    def set_triton_bundle(self, triton_bundle: Any) -> None:
+        raise NotImplementedError(type(self))
+
+
+_StrideExprStr: TypeAlias = str
+
+
+# copy_ fails when trying to write to tensors with memory overlap,
+# for expanded dimensions (a dimension which used to have size 1 -> ?)
+# we can select one element from that dimension and write to it
+# to achieve writing to all values of that dimension of the input tensor
+def get_expanded_dims(t: torch.Tensor) -> list[int]:
+    if not isinstance(t, torch.Tensor):
+        # pyrefly: ignore [bad-return]
+        return None
+    return [i for i in range(t.ndim) if t.stride(i) == 0 and t.size(i) != 1]
+
+
+def index_expanded_dims(t: torch.Tensor, expanded_dims: list[int]) -> torch.Tensor:
+    for expanded_dim in expanded_dims:
+        t = torch.ops.aten.slice(t, expanded_dim, 0, 1)
+    return t
+
+
+def complex_memory_overlap(t: torch.Tensor) -> bool:
+    if config.always_complex_memory_overlap_TESTING_ONLY:
+        return True
+
+    # if torch._debug_has_internal_overlap thinks this tensor potentially has
+    # memory overlap internally, let's dig deeper to find out whether it's true.
+    #
+    # Call squeeze() so that dimension with size 1 does not cause false positive.
+    t = index_expanded_dims(t, get_expanded_dims(t)).squeeze()
+    if torch._debug_has_internal_overlap(t) != 0:
+        strides = t.stride()
+        sizes = t.shape
+        indices = list(range(len(strides)))
+        indices = [x for _, x in sorted(zip(strides, indices))]
+        for i in range(len(strides)):
+            prev_stride = 1 if i == 0 else strides[indices[i - 1]]
+            prev_size = 1 if i == 0 else sizes[indices[i - 1]]
+            if strides[indices[i]] < prev_stride * prev_size:
+                return True
+    return False
+
+
+def maybe_handle_backward_generation(
+    compiled_graph: CompiledFxGraph,
+    boxed_forward_device_index: BoxedDeviceIndex | None,
+) -> None:
+    assert compiled_graph.current_callable is not None
+    is_backward = compiled_graph.fx_kwargs["is_backward"]
+
+    # See [Backward Generation Handling]
+    # if cudagraph'd the forward and set the device, we need to let the cudagraph manager
+    # know we are we running the backward even if we will not run it in cudagraphs
+    if is_backward and config.triton.cudagraph_trees:
+        assert boxed_forward_device_index is not None
+        assert boxed_forward_device_index.value is not None
+        compiled_graph_callable = compiled_graph.current_callable
+
+        manager = torch._inductor.cudagraph_trees.get_manager(
+            boxed_forward_device_index.value, create_if_none_exists=False
+        )
+        # should already exist from forward
+        assert manager is not None
+
+        def compiled_artifact(new_inputs: list[Any]) -> Callable[..., Any]:
+            manager.set_to_running_backward()  # type: ignore[union-attr]
+            return compiled_graph_callable(new_inputs)
+
+        compiled_graph.current_callable = compiled_artifact
+
+
+def prepare_cudagraph_post_compile(
+    compiled_graph: CompiledFxGraph,
+    example_inputs: Sequence[InputType],
+    boxed_forward_device_index: BoxedDeviceIndex | None,
+) -> None:
+    if not config.triton.cudagraph_trees:
+        # Force specialize all inputs so that CUDA graphs will work
+        for t in example_inputs:
+            if isinstance(t, torch.SymInt):
+                int(t)  # guard
+
+    is_inference = compiled_graph.fx_kwargs["is_inference"]
+    is_backward = compiled_graph.fx_kwargs["is_backward"]
+    if boxed_forward_device_index is not None and not is_inference and not is_backward:
+        boxed_forward_device_index.set(next(iter(compiled_graph.device_idxs)))
+
+
+def cudagraph_post_compile(
+    example_inputs: Sequence[InputType],
+    compiled_graph: CompiledFxGraph,
+    cudagraphs: BoxedBool,
+    constants: dict[str, torch.Tensor | type],
+    boxed_forward_device_index: BoxedDeviceIndex | None,
+) -> None:
+    """
+    Checks for any reasons not to run cudagraphs and then
+    runs it on compiled_graph.
+    Mutates the `compiled_graph.current_callable` and `cudagraphs`
+    """
+    from torch._inductor.compiler_bisector import CompilerBisector
+
+    assert compiled_graph.current_callable is not None
+    assert compiled_graph.cudagraph_info is not None
+    cached_info = compiled_graph.cudagraph_info
+    cudagraph_fail_reasons = cached_info.cudagraph_fail_reasons
+    is_inference = compiled_graph.fx_kwargs["is_inference"]
+    is_backward = compiled_graph.fx_kwargs["is_backward"]
+
+    # Check if bisector wants to disable cudagraphs for this graph
+    if CompilerBisector.disable_subsystem("inductor", "cudagraphs"):
+        BoxedBool.disable(cudagraphs)
+        maybe_handle_backward_generation(compiled_graph, boxed_forward_device_index)
+        log_cudagraph_skip_and_bump_counter("skipping cudagraphs due to bisector")
+        return
+
+    if not cudagraph_fail_reasons:
+        fx_kwargs = compiled_graph.fx_kwargs
+        static_input_idxs = fx_kwargs["static_input_idxs"]
+
+        placeholders = cached_info.placeholders
+        stack_traces = cached_info.stack_traces
+        assert stack_traces is not None, (
+            "stack_traces should not be None in cudagraph_post_compile"
+        )
+
+        prepare_cudagraph_post_compile(
+            compiled_graph, example_inputs, boxed_forward_device_index
+        )
+
+        current_callable = compiled_graph.current_callable
+        assert current_callable is not None
+        # Filter to only tensor constants (exclude opaque value type classes)
+        tensor_constants = {
+            k: v for k, v in constants.items() if isinstance(v, torch.Tensor)
+        }
+
+        device_index = next(iter(compiled_graph.device_idxs))
+        cudagraphify_kwargs = dict(
+            device_index=device_index,
+            stack_traces=stack_traces,
+            is_backward=is_backward,
+            is_inference=is_inference,
+            constants=tuple(tensor_constants.values()),
+            placeholders=placeholders,
+            mutated_input_idxs=tuple(compiled_graph.mutated_input_idxs),
+        )
+
+        policy = config.cudagraph_policy
+        if policy is not None:
+            compiled_graph.current_callable = policy.cudagraphify(
+                current_callable,
+                example_inputs,
+                static_input_idxs or (),
+                **cudagraphify_kwargs,
+            )
+        else:
+            from .compile_fx import cudagraphify
+
+            compiled_graph.current_callable = cudagraphify(
+                current_callable,
+                static_input_idxs=static_input_idxs or (),
+                **cudagraphify_kwargs,
+            )
+
+    else:
+        BoxedBool.disable(cudagraphs)
+        maybe_handle_backward_generation(compiled_graph, boxed_forward_device_index)
+
+        if "cuda" in compiled_graph.device_types:
+            # prefer better disable_cudagraphs_reason bc stack trace
+            # TODO: migrate all disable reasons to stack trace, refactor
+            if compiled_graph.disabled_cudagraphs_reason:
+                log_cudagraph_skip_and_bump_counter(
+                    compiled_graph.disabled_cudagraphs_reason
+                )
+            else:
+                log_cudagraph_skip_and_bump_counter(
+                    f"skipping cudagraphs due to {cudagraph_fail_reasons}"
+                )
+
+
+def cudagraph_partition_post_compile(
+    example_inputs: Sequence[InputType],
+    compiled_graph: CompiledFxGraph,
+    cudagraphs: BoxedBool,
+    constants: dict[str, torch.Tensor | type],
+    boxed_forward_device_index: BoxedDeviceIndex | None,
+) -> None:
+    """
+    Cudagraphify each partition functions, which first prepares the necessary
+    metadata and then applies the cudagraphify function to each partition.
+
+    Assuming all partition functions are cudagraphified and share the same order
+    as `compiled_graph.partition_maps`. See [Note: Graph Partition Map for CUDAGraph].
+    """
+    from torch._inductor.compiler_bisector import CompilerBisector
+
+    if CompilerBisector.disable_subsystem("inductor", "cudagraphs"):
+        BoxedBool.disable(cudagraphs)
+        maybe_handle_backward_generation(compiled_graph, boxed_forward_device_index)
+        log_cudagraph_skip_and_bump_counter("skipping cudagraphs due to bisector")
+        return
+
+    assert compiled_graph.cudagraph_info is not None
+    cudagraph_fail_reasons = compiled_graph.cudagraph_info.cudagraph_fail_reasons
+
+    if (
+        cudagraph_fail_reasons
+        or compiled_graph.partition_maps is None
+        or len(compiled_graph.partition_maps) == 0
+    ):
+        # cudagraphify is not called if there are no partitions
+        BoxedBool.disable(cudagraphs)
+        maybe_handle_backward_generation(compiled_graph, boxed_forward_device_index)
+        return
+
+    assert compiled_graph.current_callable is not None
+    assert compiled_graph.recursively_apply_fns is not None
+    is_inference = compiled_graph.fx_kwargs["is_inference"]
+    is_backward = compiled_graph.fx_kwargs["is_backward"]
+    static_input_idxs = OrderedSet(compiled_graph.fx_kwargs["static_input_idxs"] or ())
+    mutated_input_idxs = compiled_graph.mutated_input_idxs
+    device_index = next(iter(compiled_graph.device_idxs))
+
+    # Filter to only tensor constants (exclude opaque value type classes)
+    tensor_constants = {
+        k: v for k, v in constants.items() if isinstance(v, torch.Tensor)
+    }
+
+    assert compiled_graph.cudagraph_info.stack_traces is not None, (
+        "stack_traces should not be None in cudagraph_partition_post_compile"
+    )
+    graph_metadata = CudagraphMetadata(
+        compiled_graph.cudagraph_info.placeholders,
+        static_input_idxs,
+        mutated_input_idxs,
+        compiled_graph.cudagraph_info.stack_traces,
+        tensor_constants,
+    )
+
+    prepare_cudagraph_post_compile(
+        compiled_graph, example_inputs, boxed_forward_device_index
+    )
+
+    from .compile_fx import cudagraphify
+
+    # cudagraphify each partition function, assuming every graph partition function
+    # is cudagraphable. Non-cudagraphable ops (e.g., cpu ops) are inlined into
+    # `call` function and not included in partition functions.
+    cudagraphify_fns = []
+    for partition_map in compiled_graph.partition_maps:
+        partition_metadata = get_partition_cudagraph_metadata(
+            partition_map,
+            graph_metadata,
+        )
+
+        cudagraphify_fn = partial(
+            cudagraphify,
+            static_input_idxs=tuple(partition_metadata.static_input_idxs),
+            device_index=device_index,
+            stack_traces=partition_metadata.stack_traces,
+            is_backward=is_backward,
+            is_inference=is_inference,
+            constants=tuple(partition_metadata.constants.values()),
+            placeholders=partition_metadata.placeholders,
+            mutated_input_idxs=tuple(partition_metadata.mutated_input_idxs),
+        )
+        cudagraphify_fns.append(cudagraphify_fn)
+
+    compiled_graph.recursively_apply_fns(cudagraphify_fns)
+
+
+def maybe_realign_inputs(
+    ran_cudagraphs: BoxedBool,
+    compiled_graph: CompiledFxGraph,
+    inputs_to_check: Sequence[int],
+    mutated_inputs_idxs: OrderedSet[int],
+) -> None:
+    """
+    Realigns input strides from inputs_to_check if
+    we didn't end up running cudagraphs. Mutates
+    `compiled_graph.current_callable` if cudagraphs
+    was run. Otherwise, does nothing.
+
+    Non-mutated inputs are handled by deferred alignment copies
+    in the generated code. Only mutated inputs need the wrapper
+    for writeback.
+    """
+    if not ran_cudagraphs:
+        check_idxs = inputs_to_check
+        if compiled_graph._defers_input_alignment:
+            # Non-mutated inputs are handled by deferred alignment copies
+            # in the generated Python code. Only mutated inputs need the wrapper
+            # for writeback. Backends that don't emit deferred copies (cpp_wrapper,
+            # FXIR) need the full wrapper.
+            check_idxs = [i for i in inputs_to_check if i in mutated_inputs_idxs]
+        if check_idxs:
+            assert compiled_graph.current_callable is not None
+            new_callable = align_inputs_from_check_idxs(
+                compiled_graph.current_callable,
+                check_idxs,
+                mutated_inputs_idxs,
+            )
+            if new_callable is not compiled_graph.current_callable:
+                compiled_graph.current_callable = new_callable
+
+
+class CompiledFxGraphConstants:
+    """Wrapper class that unwraps constants from a compiled fx graph. This
+    version of the class only supports directly grabbing the saved constants off of
+    a CompiledFxGraph.
+
+    With freezing, FxGraphCache doesn't store the constants of the input
+    GraphModule it gets from AOTAutograd. Instead, it saves just the **names**
+    of those constants, and grabs the constant values directly from the graph module
+    passed in at runtime.
+
+    Thing is, we don't always *have* the graph module available at runtime, hence
+    the existence of this class and its CompiledFxGraphConstantsWithGm counterpart.
+
+    To support freezing, FXGraphCache gets passed a CompiledFxGraphConstantsWithGm during
+    post compile. Otherwise, CompiledFxGraphConstants supports the basic case of loading
+    the value of constants directly off of the original saved object.
+    """
+
+    def unwrap(self, g: CompiledFxGraph) -> dict[str, torch.Tensor | type]:
+        assert g.constants is not None
+        return {**g.constants, **g.opaque_value_type_classes}
+
+
+class CompiledFxGraphConstantsWithGm(CompiledFxGraphConstants):
+    """
+    This version of CompiledFxGraphConstants, instead of grabbing constants
+    directly saved on CompiledFxGraphs, will just grab their names. Then, it takes
+    a second GraphModule to grab the corresponding constant values out of.
+
+    This is necessary for supporting freezing in FxGraphCache.
+    """
+
+    def __init__(self, gm: torch.fx.GraphModule) -> None:
+        self.gm = gm
+
+    def unwrap(self, g: CompiledFxGraph) -> dict[str, torch.Tensor | type]:
+        frozen_params = {
+            name: getattr(self.gm, orig_name)
+            for name, orig_name in g.frozen_param_names.items()
+        }
+        constants = g.constants or {}
+        return {**constants, **frozen_params, **g.opaque_value_type_classes}
+
+
+@dataclasses.dataclass
+class CompiledFxGraph(OutputCode):
+    """
+    Class holding a compiled FX graph. This is the object serialized on disk
+    to support FxGraph caching.
+    """
+
+    current_callable: Callable[..., Any] | None
+    recursively_apply_fns: Callable[..., Any] | None
+    compiled_fn_runner: Any | None
+    cache_key: str
+    source_code: str = dataclasses.field(repr=False)  # Do not display source_code
+    runnable_graph_str: str = dataclasses.field(repr=False)  # Do not display graph
+    inductor_post_grad_graph_str: str = dataclasses.field(
+        repr=False
+    )  # Do not display graph
+    cache_linemap: list[tuple[int, str]] | None
+    device_types: OrderedSet[str]
+    device_idxs: OrderedSet[int]
+    mutated_inputs: OrderedSet[str]
+    mutated_input_idxs: OrderedSet[int]
+    constants: dict[str, torch.Tensor] | None
+    frozen_param_names: dict[str, str]
+    torchbind_constants: dict[str, torch._C.ScriptObject | FakeScriptObject]
+    opaque_value_type_classes: dict[str, type]
+    output_strides: list[tuple[_StrideExprStr, ...] | None] | None
+    disabled_cudagraphs_reason: str | None
+    metrics_deltas: metrics.CachedMetricsDeltas
+    counter_deltas: Counter[str]
+    # This is a string representation of an expression we serialize
+    # with the object so the guards can be evaluated in a different
+    # context in order to verify the validity of serving a cached
+    # fx graph. The expression must be generated by:
+    # ShapeEnv.produce_guards_expression()
+    guards_expr: str | None
+    extern_libs_key: str | None
+    inductor_provenance_mapping_str: str | None
+    inductor_provenance_stack_traces_str: str | None
+
+    cudagraph_info: CudagraphCachedInfo | None
+    partition_maps: list[GraphPartitionMap] | None
+    compile_region_name: str | None
+    fx_kwargs: _CompileFxKwargs
+    inputs_to_check: Sequence[int]
+
+    _boxed_call: bool | None = None
+    _triton_bundle: TritonBundle | None = None
+    _wrap_compiled_regions: bool = False
+    _defers_input_alignment: bool = False
+    # Metadata-stripped copy of the FX graph for fake tensor propagation.
+    # Running this graph under FakeTensorMode re-derives output shapes
+    # (including aliasing) from the input shapes.
+    _original_gm: torch.fx.GraphModule | None = None
+
+    def __init__(
+        self,
+        current_callable: Callable[..., Any] | None,
+        graph: GraphLowering,
+        gm: torch.fx.GraphModule,
+        output_strides: list[tuple[_StrideExprStr, ...] | None],
+        disabled_cudagraphs_reason: str | None,
+        metrics_deltas: metrics.CachedMetricsDeltas,
+        counter_deltas: Counter[str],
+        cudagraphs: BoxedBool,
+        example_inputs: Sequence[InputType],
+        static_input_idxs: Sequence[int],
+        compile_region_name: str | None,
+        fx_kwargs: _CompileFxKwargs,
+        inputs_to_check: Sequence[int],
+        runnable_graph_str: str,
+        inductor_post_grad_graph_str: str,
+        compiled_fn_runner: Any | None = None,
+        inductor_provenance_mapping_str: str | None = None,
+        inductor_provenance_stack_traces_str: str | None = None,
+    ) -> None:
+        self.current_callable = current_callable
+        self.compiled_fn_runner = compiled_fn_runner
+        self.recursively_apply_fns = (
+            compiled_fn_runner.recursively_apply_fns
+            if compiled_fn_runner is not None
+            else None
+        )
+        self.cache_key = graph.cache_key
+        if graph.cache_path:
+            with open(graph.cache_path) as f:
+                self.source_code = f.read()
+        self.runnable_graph_str = runnable_graph_str
+        self.inductor_post_grad_graph_str = inductor_post_grad_graph_str
+        self.inductor_provenance_mapping_str = inductor_provenance_mapping_str
+        self.inductor_provenance_stack_traces_str = inductor_provenance_stack_traces_str
+        self.cache_linemap = graph.cache_linemap
+        # TODO - ordered set
+        self.device_types = OrderedSet(graph.device_types)
+        self.device_idxs = OrderedSet(graph.device_idxs)
+        self.mutated_inputs = OrderedSet(graph.mutated_inputs)
+        self.mutated_input_idxs = OrderedSet(graph.mutated_input_idxs)
+
+        # We store the constant attributes in the cache entry and re-attach them
+        # to the module created in PyCodeCache.load_by_key_path. In the case that
+        # the graph has frozen parameters, we save the mapping from the attribute
+        # names in the GraphLowering to the original name of the attribute in the
+        # GraphModule. When we create the module from the cache entry, we then
+        # look up the constants from the current GraphModule. This scheme allows
+        # us to support caching with freezing.
+        if not has_frozen_params(gm):
+            self.constants = graph.constants
+            self.frozen_param_names = {}
+        else:
+            self.constants = {}
+            self.frozen_param_names = {}
+            for k, v in graph.constants.items():
+                if is_frozen_param(v):
+                    self.frozen_param_names[k] = graph.allocated_constant_name[k]
+                else:
+                    self.constants[k] = v
+
+        self.torchbind_constants = graph.torchbind_constants
+        self.opaque_value_type_classes = graph.opaque_value_type_classes
+        self.output_strides = output_strides
+        self.disabled_cudagraphs_reason = disabled_cudagraphs_reason
+        self.metrics_deltas = metrics_deltas
+        self.counter_deltas = counter_deltas
+        self.guards_expr = None
+        self.extern_libs_key = None
+        self.cudagraph_info = None
+        self.partition_maps = graph.partition_maps
+        self._defers_input_alignment = getattr(graph, "_defers_input_alignment", False)
+        self.fx_kwargs = {}
+        self.inputs_to_check = ()
+
+        cudagraph_info = None
+        if cudagraphs:
+            # check cudagraph disabling reasons from inductor lowering
+            if self.disabled_cudagraphs_reason:
+                if "cuda" in self.device_types:
+                    log_cudagraph_skip_and_bump_counter(
+                        f"skipping cudagraphs due to {self.disabled_cudagraphs_reason}"
+                    )
+                else:
+                    counters["inductor"]["cudagraph_skips"] += 1
+                BoxedBool.disable(cudagraphs)
+            else:
+                complex_memory_overlap_inputs = any(
+                    complex_memory_overlap(t)
+                    for t in example_inputs
+                    if isinstance(t, torch.Tensor)
+                )
+
+                if not config.triton.cudagraph_support_input_mutation:
+                    # Skip supports for cudagraph-managed tensors
+                    from torch._inductor.cudagraph_utils import (
+                        check_for_mutation_ignore_cuda_graph_managed_tensor,
+                    )
+
+                    has_mutation_str = (
+                        check_for_mutation_ignore_cuda_graph_managed_tensor(
+                            gm,
+                            self.mutated_inputs,
+                            self.mutated_input_idxs,
+                            static_input_idxs,
+                        )
+                    )
+                    has_mutation = has_mutation_str is not None
+
+                    if has_mutation:
+                        self.disabled_cudagraphs_reason = has_mutation_str
+                else:
+                    # Check mutation later to support cudagraph-managed tensors
+                    has_mutation = None
+
+                cudagraph_tests = [
+                    (not has_mutation, "mutated inputs"),
+                    (not complex_memory_overlap_inputs, "complex memory overlap"),
+                    (
+                        all(
+                            isinstance(
+                                t,
+                                (
+                                    torch.Tensor,
+                                    torch.SymInt,
+                                    torch.Generator,
+                                    OpaqueBase,
+                                ),
+                            )
+                            for t in example_inputs
+                        ),
+                        "non-Tensor inputs",
+                    ),
+                ]
+                output = output_node(gm)
+                # output args are tuple of first argument
+                assert len(output.args) == 1
+                # Use stack traces captured on the output node before
+                # post-grad passes, which may strip stack_trace from
+                # individual arg nodes.
+                stack_traces = output.meta.get("output_stack_traces") or [
+                    (arg.stack_trace if isinstance(arg, torch.fx.node.Node) else None)
+                    for arg in output.args[0]  # type: ignore[union-attr]
+                ]
+                cudagraph_fail_reasons = [s for b, s in cudagraph_tests if not b]
+                placeholders = tuple(get_placeholder_info(gm.graph))
+                cudagraph_info = CudagraphCachedInfo(
+                    placeholders, stack_traces, cudagraph_fail_reasons
+                )
+
+        self.cudagraph_info = cudagraph_info
+        self.compile_region_name = compile_region_name
+        self.inputs_to_check = inputs_to_check
+        self.fx_kwargs = fx_kwargs
+
+        # aot autograd needs to know to pass in inputs as a list
+        self._boxed_call = True
+
+        # Store whether to wrap compiled regions in inductor_compiled_code HOP
+        # This is set at compile time to avoid runtime overhead
+        self._wrap_compiled_regions = config.wrap_inductor_compiled_regions
+
+        if self._wrap_compiled_regions:
+            # Store a metadata-stripped copy of the FX graph. Running this
+            # under FakeTensorMode re-derives output shapes and aliasing
+            # from the input fake tensors.
+            import copy
+
+            gm_copy = copy.deepcopy(gm)
+            for node in gm_copy.graph.nodes:
+                node.meta.clear()
+            self._original_gm = gm_copy
+
+    def __del__(self) -> None:
+        if self.compiled_fn_runner is not None:
+            # For torch._inductor.config.graph_partition = True,
+            # self.compiled_fn_runner.partitions hold cudagraphified functions
+            # which prevents deallocation. When CompiledFxGraph is deleted,
+            # self.compiled_fn_runner will not be called in the future so we
+            # should also delete these partitions.
+            del self.compiled_fn_runner.partitions
+
+    def __call__(self, inputs: Sequence[Any]) -> Any:
+        assert self.current_callable is not None
+
+        if (
+            torch._inductor.debug.RECORD_GRAPH_EXECUTION
+            and torch._inductor.debug.GRAPH_EXECUTION_ORDER is not None
+        ):
+            graph_id = self.fx_kwargs.get("graph_id")
+            compile_id = (
+                torch._inductor.debug.GRAPH_COMPILE_IDS.get(graph_id)
+                if graph_id is not None
+                and torch._inductor.debug.GRAPH_COMPILE_IDS is not None
+                else None
+            )
+            torch._inductor.debug.GRAPH_EXECUTION_ORDER.append(
+                {
+                    "compile_id": compile_id,
+                }
+            )
+        try:
+            # Checking the profiler directly is faster than nullcontext
+            if torch.autograd.profiler._is_profiler_enabled:
+                with torch._C._profiler._RecordFunctionFast(
+                    f"## Call CompiledFxGraph {self._fx_graph_cache_key} ##",
+                    keyword_values={"scope": "user_scope"},
+                ):
+                    return self.current_callable(inputs)
+            else:
+                return self.current_callable(inputs)
+        finally:
+            get_runtime_metrics_context().finish()
+            AutotuneCacheBundler.end_compile()
+
+    def post_compile(
+        self,
+        example_inputs: Sequence[InputType],
+        constants: CompiledFxGraphConstants,
+        graph_kwargs: _CompileFxKwargs,
+    ) -> None:
+        """
+        Run a set of post processing steps after loading from the cache. These involve:
+         - Setting the tracing context output strides
+         - Running cudagraphs if enabled
+         - Realigning inputs
+
+        This runs whether or not we have a cache hit, and always runs directly after we get a CompiledFxGraph.
+        The results of this function are *not* saved in the cache itself.
+        """
+        if config.graph_partition and _unstable_customized_partition_wrapper.wrapper:
+            # Mechanically apply user-specified cudagraph wrappers without modification
+            assert self.recursively_apply_fns is not None
+            assert self.compiled_fn_runner is not None
+            num_partitions = len(self.compiled_fn_runner.partitions)
+            wrapper_metadatas = [
+                CUDAGraphWrapperMetadata(num_partitions, i)
+                for i in range(num_partitions)
+            ]
+            customized_wrapper = _unstable_customized_partition_wrapper.wrapper
+            customized_wrappers_with_metadata = [
+                lambda f, m=metadata: customized_wrapper(f, m)
+                for metadata in wrapper_metadatas
+            ]
+            self.recursively_apply_fns(customized_wrappers_with_metadata)
+            return
+
+        set_tracing_context_output_strides(example_inputs, self)
+        assert graph_kwargs["cudagraphs"] is not None
+        assert graph_kwargs["is_backward"] is not None
+        is_backward = graph_kwargs["is_backward"]
+        cudagraphs: BoxedBool = graph_kwargs["cudagraphs"]
+
+        # When a CUDAGraphPolicy is set and it says not to wrap this
+        # inner CompiledFxGraph (e.g. because wrapping happens at the
+        # outer level via policy.wrap_output), disable cudagraphs for
+        # this graph so the rest of post_compile (input realignment,
+        # _wrap_compiled_regions) still runs normally.
+        policy = config.cudagraph_policy
+        if policy is not None and not policy.should_wrap(self):
+            counters["inductor"]["cudagraph_skips"] += 1
+            BoxedBool.disable(cudagraphs)
+
+        if cudagraphs:
+            # It's possible that cudagraphs is enabled, but was disabled
+            # during a previous compilation we're loading from the cache.
+            # If so, we need to disable it on this new process too.
+            if self.disabled_cudagraphs_reason:
+                if "cuda" in self.device_types:
+                    log_cudagraph_skip_and_bump_counter(
+                        f"skipping cudagraphs due to {self.disabled_cudagraphs_reason}"
+                    )
+                else:
+                    counters["inductor"]["cudagraph_skips"] += 1
+                BoxedBool.disable(cudagraphs)
+            else:
+                if is_backward:
+                    assert "boxed_forward_device_index" in graph_kwargs
+                    boxed_forward_device_index = graph_kwargs[
+                        "boxed_forward_device_index"
+                    ]
+                else:
+                    # On the forward we don't know whether or not
+                    # boxed_forward_device_index is set yet
+                    boxed_forward_device_index = graph_kwargs.get(
+                        "boxed_forward_device_index", None
+                    )
+
+                if config.graph_partition and policy is None:
+                    # With graph_partition=True, we skip some cudagraph checks
+                    # if it's supported with partition, so we use
+                    # cudagraph_partition_post_compile.  When a CUDAGraphPolicy
+                    # is active, we use cudagraph_post_compile instead so the
+                    # policy controls wrapping via policy.cudagraphify().
+                    cudagraph_partition_post_compile(
+                        example_inputs,
+                        self,
+                        cudagraphs,
+                        constants.unwrap(self),
+                        boxed_forward_device_index,
+                    )
+                else:
+                    cudagraph_post_compile(
+                        example_inputs,
+                        self,
+                        cudagraphs,
+                        constants.unwrap(self),
+                        boxed_forward_device_index,
+                    )
+        inputs_to_check = self.inputs_to_check
+        # cudagraphs could have been disabled from the earlier conditions
+        # so we still need to realign inputs if that happens
+        maybe_realign_inputs(
+            cudagraphs,
+            self,
+            inputs_to_check,
+            self.mutated_input_idxs,
+        )
+
+        # Apply inductor_compiled_code HOP wrapper if configured
+        # This is done in post_compile to ensure it works with cached artifacts
+        if self._wrap_compiled_regions and self.current_callable is not None:
+            from torch._higher_order_ops.wrap import InductorCompiledCallable
+
+            original_callable = self.current_callable
+
+            inductor_callable = InductorCompiledCallable(
+                original_callable,
+                self._original_gm,
+                compile_region_name=self.compile_region_name,
+            )
+
+            def wrapped_callable(inputs):
+                if is_in_torch_dispatch_mode():
+                    kwargs = (
+                        {"name": self.compile_region_name}
+                        if self.compile_region_name is not None
+                        else {}
+                    )
+                    return inductor_compiled_code(inductor_callable, inputs, **kwargs)
+                else:
+                    return original_callable(inputs)
+
+            self.current_callable = wrapped_callable
+
+    def set_triton_bundle(self, triton_bundle: Any) -> None:
+        self._triton_bundle = triton_bundle
+
+    def prepare_for_serialization(self) -> None:
+        # We can't really serialize callables that may be C++/Triton/etc.,
+        # so we serialize their PyCodeCache disk cache location instead.
+        # TODO: This could be better if we're ever able to serialize compiled
+        # models to disk.
+        self.current_callable = None
+        self.recursively_apply_fns = None
+        self.compiled_fn_runner = None
+        # Note: _original_gm is already picklable (metadata stripped at creation)
+        # Note: _serialized_fx_graph is already in serializable form (SerializedGraphModule)
+        # so it doesn't need to be cleared
+
+    def write_to_disk(self) -> str:
+        from torch._dynamo.utils import counters
+        from torch._inductor.codecache import get_path, write_atomic
+
+        # See _save_graph(); we don't store the callable in the cache entry so
+        # recreate it here from the PyCodeCache disk cache.
+        artifact_path = get_path(self.cache_key, "py")[2]
+        code = self.source_code
+        if not os.path.exists(artifact_path):
+            counters["inductor"]["fxgraph_lookup_write_file"] += 1
+            write_atomic(artifact_path, code, make_dirs=True)
+        return artifact_path
+
+    def after_deserialization(self, constants: CompiledFxGraphConstants) -> str:
+        from torch._dynamo.utils import dynamo_timed
+        from torch._inductor.codecache import PyCodeCache
+
+        artifact_path = self.write_to_disk()
+
+        try:
+            with dynamo_timed(
+                "PyCodeCache.load_by_key_path",
+                log_pt2_compile_event=True,
+            ):
+                code_cache = PyCodeCache.load_by_key_path(
+                    self.cache_key,
+                    artifact_path,
+                    self.cache_linemap,
+                    constants.unwrap(self),
+                )
+                self.current_callable = code_cache.call
+                self.recursively_apply_fns = getattr(
+                    code_cache, "recursively_apply_fns", None
+                )
+                self.compiled_fn_runner = getattr(code_cache, "runner", None)
+        except OSError:
+            log.error("Failed to load artifact: %s", artifact_path)
+            raise
+
+        return artifact_path
+
+
+@dataclasses.dataclass
+class CompiledAOTI(OutputCode):
+    """
+    Class holding an AOTInductor compiled so.
+    """
+
+    filename: str | list[str | Weights] | torch.fx.GraphModule
+    device_type: str
+    current_callable: Callable[..., Any] | None = None
+    _cached_files: dict[str, bytes] = dataclasses.field(default_factory=dict)
+
+    def __post_init__(self):
+        if not config.aot_inductor.link_libtorch:
+            return
+
+        if (
+            torch._inductor.cpp_builder._IS_MACOS
+            or torch._inductor.cpp_builder._IS_WINDOWS
+        ):
+            return
+
+        if config.aot_inductor.cross_target_platform == "windows":
+            return
+
+        if config.aot_inductor.package_cpp_only:
+            return
+
+        if not config.enable_autograd_for_aot:
+            return
+
+        if isinstance(self.filename, list):
+            current_callable = next(
+                fn for fn in self.filename if isinstance(fn, str) and fn.endswith(".so")
+            )
+        else:
+            current_callable = self.filename
+
+        if isinstance(current_callable, torch.fx.GraphModule):
+            # pyrefly: ignore [bad-assignment]
+            self.current_callable = current_callable
+            return
+
+        if self.device_type.startswith("cuda"):
+            current_callable = (
+                torch._C._aoti.AOTIModelContainerRunnerCuda(  # type: ignore[call-arg]
+                    current_callable,
+                    1,
+                    self.device_type,
+                    "",
+                    True,
+                ).run  # type: ignore[attr-defined]
+            )  # type: ignore[attr-defined]
+        elif self.device_type.startswith("xpu"):
+            current_callable = (
+                torch._C._aoti.AOTIModelContainerRunnerXpu(  # type: ignore[call-arg]
+                    current_callable,
+                    1,
+                    self.device_type,
+                    "",
+                ).run  # type: ignore[attr-defined]
+            )  # type: ignore[attr-defined]
+        elif self.device_type == "cpu":
+            current_callable = (
+                torch._C._aoti.AOTIModelContainerRunnerCpu(  # type: ignore[call-arg]
+                    current_callable, 1
+                ).run  # type: ignore[attr-defined]
+            )  # type: ignore[attr-defined]
+        else:
+            raise RuntimeError(f"unsupported device type {self.device_type}")
+        self.current_callable = current_callable
+        self._boxed_call = True
+        for file in self._cached_files:
+            if not os.path.exists(file):
+                with open(file, "wb") as f:
+                    f.write(self._cached_files[file])
+
+    def __call__(self, inputs: Sequence[Any]) -> Any:
+        if self.current_callable is None:
+            raise RuntimeError("AOTInductor compiled so is not loaded")
+        return self.current_callable(inputs)
+
+    def prepare_for_serialization(self) -> None:
+        self.current_callable = None
+        self._cached_files = {}
+        filenames: list[str] = []
+        if isinstance(self.filename, list):
+            filenames = self.filename  # type: ignore[assignment]
+        elif isinstance(self.filename, str):
+            filenames = [self.filename]
+        for name in filenames:
+            with open(name, "rb") as f:
+                self._cached_files[name] = f.read()
+
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        state["current_callable"] = None
+        return state
+
+    def post_compile(
+        self,
+        example_inputs: Sequence[InputType],
+        constants: CompiledFxGraphConstants,
+        graph_kwargs: _CompileFxKwargs,
+    ) -> None:
+        if self.current_callable is None:
+            self.__post_init__()
+
+    def set_triton_bundle(self, triton_bundle: Any) -> None:
+        pass
+
+
+@dataclasses.dataclass
+class MockFXGraphCacheOutput(OutputCode):
+    gm: Any = None
+
+    def __post_init__(self) -> None:
+        self._boxed_call = True
+
+    def post_compile(
+        self,
+        example_inputs: Sequence[InputType],
+        constants: CompiledFxGraphConstants,
+        graph_kwargs: _CompileFxKwargs,
+    ) -> None:
+        pass
+
+    def __call__(self, inputs: Sequence[Any]) -> Any:
+        return self.gm(inputs)
+
+    def set_triton_bundle(self, triton_bundle: Any) -> None:
+        pass
+
+
+@dataclasses.dataclass
+class RegionalOutputCode(OutputCode):
+    """
+    OutputCode for regional inductor compilation results.
+
+    Regional inductor returns a torch.fx.GraphModule that contains both
+    compiled regions (via standalone_compile) and eager regions. This needs
+    special serialization using GraphPickler instead of standard pickle.
+
+    The serialization strategy stores the GraphModule as bytes using
+    GraphPickler.dumps(), which handles FakeTensors, AOTCompiledArtifacts,
+    and other special objects that standard pickle cannot handle.
+    """
+
+    # The serialized graph module as bytes (using GraphPickler)
+    _serialized_graph_module: bytes | None = None
+
+    # The actual graph module (cleared during serialization)
+    _graph_module: torch.nn.Module | None = None
+
+    # Optional filter for ops during serialization
+    _ops_filter: Callable[[str], bool] | None = None
+
+    def __init__(
+        self,
+        graph_module: torch.fx.GraphModule,
+        ops_filter: Callable[[str], bool] = _ops_filter_safe,
+        node_metadata_key_filter: Callable[[str], bool] | None = (
+            _node_metadata_key_filter_safe
+        ),
+    ):
+        """
+        Args:
+            graph_module: The torch.fx.GraphModule returned by regional_inductor
+            ops_filter: Optional filter function for op names during serialization.
+                If provided, only ops whose name passes the filter will be serialized.
+        """
+        super().__init__()
+        self._graph_module = graph_module
+        self._serialized_graph_module = None
+        self._serialized_wrappers = []
+        self._boxed_call = True
+        _, module = self._unwrap_graph_module()
+        self._inner_boxed_call = isinstance(
+            module.graph._codegen, torch.fx.graph._BoxedCodeGen
+        )
+        self._ops_filter = ops_filter
+        self._node_metadata_key_filter = node_metadata_key_filter
+
+    def __call__(self, inputs: list[Any]) -> Any:
+        """
+        Execute the regional compiled graph.
+
+        Args:
+            inputs: A mutable list of inputs. Must be a list (not an arbitrary
+                Sequence) because the boxed calling convention allows the callee
+                to mutably clear the list when inputs become dead, enabling early
+                memory deallocation.
+        """
+        if self._graph_module is None:
+            raise RuntimeError(
+                "RegionalOutputCode has no graph module loaded. "
+                "Did you forget to call post_compile()?"
+            )
+
+        if self._inner_boxed_call:
+            return self._graph_module(inputs)
+        return self._graph_module(*inputs)
+
+    @property
+    def graph(self):
+        _, module = self._unwrap_graph_module()
+        return module.graph
+
+    def _unwrap_graph_module(
+        self,
+    ) -> tuple[list[tuple[Callable, dict[str, Any]]], torch.fx.GraphModule]:
+        module = self._graph_module
+        serialized_wrappers = []
+        if isinstance(module, torch._dynamo.OptimizedModule):
+            dynamo_ctx = module.dynamo_ctx
+            assert isinstance(dynamo_ctx, torch._dynamo.eval_frame.DisableContext)
+            serialized_wrappers.append(
+                (
+                    torch._dynamo.disable,
+                    {"reason": dynamo_ctx.msg, "wrapping": dynamo_ctx.wrapping},
+                )
+            )
+            module = module._orig_mod
+        assert isinstance(module, torch.fx.GraphModule)
+        return serialized_wrappers, module
+
+    def post_compile(
+        self,
+        example_inputs: Sequence[InputType],
+        constants: CompiledFxGraphConstants,
+        graph_kwargs: _CompileFxKwargs,
+    ) -> None:
+        """
+        Post-compile processing for regional inductor.
+
+        This deserializes the GraphModule from bytes using GraphPickler,
+        extracting the fake_mode from example_inputs.
+        """
+        if self._graph_module is not None:
+            return
+        assert self._serialized_graph_module is not None
+        # Get fake mode from example inputs
+        from torch._guards import detect_fake_mode
+
+        fake_mode = detect_fake_mode(example_inputs)
+        if fake_mode is None:
+            raise RuntimeError(
+                "Could not detect fake mode from example inputs. "
+                "Regional inductor requires fake mode for deserialization."
+            )
+
+        # Deserialize the graph module
+        from torch.fx._graph_pickler import GraphPickler
+
+        gm = GraphPickler.loads(self._serialized_graph_module, fake_mode)
+        assert isinstance(gm, torch.fx.GraphModule)
+        gm.recompile()
+        for fn, kwargs in reversed(self._serialized_wrappers):
+            gm = fn(gm, **kwargs)
+        self._graph_module = gm
+
+    def set_triton_bundle(self, triton_bundle: Any) -> None:
+        """Regional inductor doesn't use triton bundles directly."""
+
+    def prepare_for_serialization(self) -> None:
+        """
+        Prepare for serialization by converting the GraphModule to bytes.
+
+        This uses GraphPickler to serialize the graph module since it contains
+        special objects like FakeTensors and AOTCompiledArtifacts that need
+        custom pickling.
+        """
+        if self._graph_module is not None:
+            from torch.fx._graph_pickler import GraphPickler, Options
+
+            self._serialized_graph_module = None
+            self._serialized_wrappers, graph_module = self._unwrap_graph_module()
+
+            self._serialized_graph_module = GraphPickler.dumps(
+                graph_module,
+                options=Options(
+                    ops_filter=self._ops_filter,
+                    node_metadata_key_filter=self._node_metadata_key_filter,
+                ),
+            )
+            # Clear the graph module to avoid pickling it with standard pickle
+            self._graph_module = None

outputs/audit_venv/lib/python3.11/site-packages/torch/_inductor/package/__init__.py

@@ -0,0 +1 @@
+from .package import AOTICompiledModel, load_package, package_aoti

outputs/audit_venv/lib/python3.11/site-packages/torch/_inductor/package/build_package.py

@@ -0,0 +1,15 @@
+build_package_contents = """
+import os
+from pathlib import Path
+
+from torch._inductor.package.package import compile_so
+
+curr_dir = Path(__file__).parent
+aoti_files = [
+    os.path.join(root, file)
+    for root, dirs, files in os.walk(curr_dir)
+    for file in files
+]
+
+output_so = compile_so(curr_dir, aoti_files, curr_dir)
+"""

outputs/audit_venv/lib/python3.11/site-packages/torch/_inductor/package/package.py

@@ -0,0 +1,138 @@
+import io
+import json
+import logging
+import os
+import tempfile
+from typing import IO
+
+import torch
+from torch._inductor import config
+from torch._inductor.cpp_builder import BuildOptionsBase, CppBuilder
+from torch.export.pt2_archive._package import (
+    AOTI_FILES,
+    AOTICompiledModel,
+    load_pt2,
+    package_pt2,
+)
+from torch.types import FileLike
+
+
+log = logging.getLogger(__name__)
+
+
+def compile_so(aoti_dir: str, aoti_files: list[str], so_path: str) -> str:
+    def get_aoti_file_with_suffix(suffix: str) -> str:
+        for file in aoti_files:
+            if file.endswith(suffix):
+                return file
+        raise RuntimeError(f"Unable to find file with suffix {suffix}")
+
+    # Compile all the files into a .so
+    cpp_file = os.path.join(aoti_dir, get_aoti_file_with_suffix(".cpp"))
+    consts_o = os.path.join(aoti_dir, get_aoti_file_with_suffix(".o"))
+
+    file_name = os.path.splitext(cpp_file)[0]
+
+    # Parse compile flags and build the .o file
+    with open(file_name + "_compile_flags.json") as f:
+        compile_flags = json.load(f)
+
+    compile_options = BuildOptionsBase(
+        **compile_flags, use_relative_path=config.is_fbcode()
+    )
+    object_builder = CppBuilder(
+        name=file_name,
+        sources=cpp_file,
+        BuildOption=compile_options,
+    )
+    output_o = object_builder.get_target_file_path()
+    object_builder.build()
+
+    # Parse linker flags and build the .so file
+    with open(file_name + "_linker_flags.json") as f:
+        linker_flags = json.load(f)
+
+    linker_options = BuildOptionsBase(
+        **linker_flags, use_relative_path=config.is_fbcode()
+    )
+    so_builder = CppBuilder(
+        name=os.path.split(so_path)[-1],
+        sources=[output_o, consts_o],
+        BuildOption=linker_options,
+        output_dir=so_path,
+    )
+    output_so = so_builder.get_target_file_path()
+    so_builder.build()
+
+    # mmapped weights
+    serialized_weights_filename = file_name + "_serialized_weights.bin"
+    if serialized_weights_filename in aoti_files:
+        with open(serialized_weights_filename, "rb") as f_weights:
+            serialized_weights = f_weights.read()
+
+        with open(output_so, "a+b") as f_so:
+            so_size = f_so.tell()
+            # Page align the weights
+            f_so.write(b" " * (16384 - so_size % 16384))
+            f_so.write(serialized_weights)
+
+    return output_so
+
+
+def package_aoti(
+    archive_file: FileLike,
+    aoti_files: AOTI_FILES,
+) -> FileLike:
+    """
+    Saves the AOTInductor generated files to the PT2Archive format.
+
+    Args:
+        archive_file: The file name to save the package to.
+        aoti_files: This can either be a singular path to a directory containing
+        the AOTInductor files, or a dictionary mapping the model name to the
+        path to its AOTInductor generated files.
+    """
+
+    return package_pt2(
+        archive_file,
+        aoti_files=aoti_files,
+    )
+
+
+def load_package(
+    path: FileLike,
+    model_name: str = "model",
+    run_single_threaded: bool = False,
+    num_runners: int = 1,
+    device_index: int = -1,
+) -> AOTICompiledModel:
+    try:
+        pt2_contents = load_pt2(
+            path,
+            run_single_threaded=run_single_threaded,
+            num_runners=num_runners,
+            device_index=device_index,
+        )
+        if model_name not in pt2_contents.aoti_runners:
+            raise RuntimeError(f"Model {model_name} not found in package")
+        return pt2_contents.aoti_runners[model_name]
+    except RuntimeError:
+        log.warning("Loading outdated pt2 file. Please regenerate your package.")
+
+    if isinstance(path, (io.IOBase, IO)):
+        with tempfile.NamedTemporaryFile(suffix=".pt2") as f:
+            # TODO(angelayi): We shouldn't need to do this -- miniz should
+            # handle reading the buffer. This is just a temporary workaround
+            path.seek(0)
+            f.write(path.read())
+            log.debug("Writing buffer to tmp file located at %s.", f.name)
+            loader = torch._C._aoti.AOTIModelPackageLoader(
+                f.name, model_name, run_single_threaded, num_runners, device_index
+            )
+            return AOTICompiledModel(loader)
+
+    path = os.fspath(path)  # AOTIModelPackageLoader expects (str, str)
+    loader = torch._C._aoti.AOTIModelPackageLoader(
+        path, model_name, run_single_threaded, num_runners, device_index
+    )
+    return AOTICompiledModel(loader)

outputs/audit_venv/lib/python3.11/site-packages/torch/_inductor/pattern_matcher.py

@@ -0,0 +1,2495 @@
+"""
+# Inductor Pattern Matcher
+
+The pattern matcher enables search/replace within an FX graph.
+
+The main entrypoint to the pattern matcher is register_replacement(). Given a
+search function and a replacement function this will register a replacement with
+a pass (such as torch._inductor.fx_passes.joint_graph.patterns).
+
+Internally the pattern matcher represents patterns as a graph (a DAG). Creating
+new patterns manually as a graph is cumbersome and error-prone so the standard
+way to create patterns (using register_replacement()) is to provide a search
+function and a replacement function which is traced and converted into a graph.
+
+Because the search functions are built somewhat generic (they tend to ignore
+tensor sizes, for example) register_replacement() allows you to specify an
+`extra_check` function which performs additional checks to verify that the
+matched pattern fully matches before returning it.
+
+## Precompiled Patterns
+
+New patterns are added using register_replacement(). Patterns added in this way
+can have a compile-time overhead because they need to be traced before
+use. Patterns can be precompiled and added using gen_register_replacement()
+instead. To do this you call gen_register_replacement() instead of
+register_replacement(). The arguments are the same except for an additional
+unique name which is used as a lookup key.
+
+## Internals
+
+The match DAG is represented by a graph of `PatternExpr` nodes. Each PatternExpr
+implements a `_match` method which returns either a `Match` object for a
+successful match or a `FailedMatch` object for a failure to match.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import dataclasses
+import functools
+import importlib
+import inspect
+import itertools
+import logging
+import operator
+import os
+import re
+import textwrap
+import typing
+from abc import ABC, abstractmethod
+from collections import defaultdict
+from collections.abc import Callable, Collection, Generator, Iterable, Mapping, Sequence
+from pathlib import Path
+from typing import Any, NoReturn, Protocol, TypeVar
+from typing_extensions import Self, TypeIs
+
+import torch
+import torch._guards
+import torch.fx
+import torch.utils._pytree as pytree
+from torch._dispatch.python import enable_python_dispatcher
+from torch._dynamo.utils import counters
+from torch._prims_common import is_integer_dtype
+from torch._subclasses.fake_tensor import unset_fake_temporarily
+from torch.fx.experimental.proxy_tensor import make_fx
+from torch.fx.experimental.symbolic_shapes import guard_or_false, statically_known_true
+from torch.fx.graph_module import _get_attr
+from torch.fx.immutable_collections import immutable_dict, immutable_list
+from torch.fx.passes.graph_transform_observer import GraphTransformObserver
+from torch.fx.traceback import preserve_node_meta
+from torch.utils._ordered_set import OrderedSet
+
+from .._functorch import config as functorch_config
+from .._functorch.aot_autograd import aot_function, make_boxed_func
+from .._functorch.partitioners import default_partition
+from .._subclasses import FakeTensor, FakeTensorMode
+from ..fx import Transformer
+from . import config
+from .decomposition import select_decomp_table
+from .lowering import fallback_node_due_to_unsupported_type
+
+
+log = logging.getLogger(__name__)
+aten = torch.ops.aten
+prims = torch.ops.prims
+
+Constant = Any
+NodeOrConstant = Constant | torch.fx.Node
+
+backend = os.environ.get("TORCHINDUCTOR_PATTERN_MATCH_BACKEND", "inductor")
+
+
+_debug_nodes_cache: bool | OrderedSet[str] | None = None
+_debug_nodes_env_value_cache: str | None = None
+
+
+def _should_debug_node(node_name: str) -> bool:
+    def _get_debug_nodes() -> bool | OrderedSet[str]:
+        global _debug_nodes_cache, _debug_nodes_env_value_cache
+
+        def parse_debug_env(env_value: str | None) -> bool | OrderedSet[str]:
+            if not env_value:
+                return False
+            if env_value == "all":
+                return True
+            return OrderedSet(env_value.split(","))
+
+        current_env = os.environ.get("TORCHINDUCTOR_PATTERN_MATCH_DEBUG")
+
+        # Recompute only if env changed
+        if current_env != _debug_nodes_env_value_cache or _debug_nodes_cache is None:
+            _debug_nodes_cache = parse_debug_env(current_env)
+            _debug_nodes_env_value_cache = current_env
+
+        return _debug_nodes_cache
+
+    debug_nodes = _get_debug_nodes()
+    if isinstance(debug_nodes, bool):
+        return debug_nodes
+    return node_name in debug_nodes
+
+
+class SearchFn(Protocol):
+    __name__: str
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any: ...
+
+
+class ReplaceFn(Protocol):
+    def __call__(self, *args: Any, **kwargs: Any) -> Any: ...
+
+
+class TraceFn(Protocol):
+    def __call__(
+        self, fn: SearchFn | ReplaceFn, *args: Any, **kwargs: Any
+    ) -> torch.fx.GraphModule: ...
+
+
+T = TypeVar("T")
+
+# What's a better name for this?
+FnsType = torch.fx.node.Target | str
+
+
+class Multiple:
+    def __init__(self) -> None:
+        # Ensure we're really a singleton.
+        assert "MULTIPLE" not in globals() or self is MULTIPLE
+
+
+# Sentinel indicating multiple quantities can be matched
+MULTIPLE = Multiple()
+
+
+def _transfer_meta(
+    new_meta: dict[str, Any], old_node: torch.fx.Node, pass_name: str = ""
+) -> None:
+    from torch.fx.traceback import NodeSource, NodeSourceAction
+
+    # Transfer metadata after pattern matching occurs.
+    # Copies _COPY_META_FIELDS, stack_trace, and (if missing) val/tensor_meta.
+    if config.trace.provenance_tracking_level == 1:
+        new_from_node = new_meta.get("from_node", []).copy()
+        new_from_node.append(NodeSource(old_node, pass_name, NodeSourceAction.REPLACE))
+        new_meta.update(
+            (k, v)
+            for k, v in old_node.meta.items()
+            if k in torch.fx.proxy._COPY_META_FIELDS
+        )
+        new_meta["from_node"] = new_from_node
+    else:
+        new_meta.update(
+            (k, v)
+            for k, v in old_node.meta.items()
+            if k in torch.fx.proxy._COPY_META_FIELDS
+        )
+    if "stack_trace" in old_node.meta:
+        new_meta["stack_trace"] = old_node.meta["stack_trace"]
+    # Copy val/tensor_meta only when the new node doesn't already have them
+    # (e.g. from tracing the replacement graph). Don't overwrite if present
+    # since the replacement's own val is more accurate.
+    if "val" not in new_meta and "val" in old_node.meta:
+        new_meta["val"] = old_node.meta["val"]
+    if "tensor_meta" not in new_meta and "tensor_meta" in old_node.meta:
+        new_meta["tensor_meta"] = old_node.meta["tensor_meta"]
+
+
+class Match:
+    """
+    Represents a successfully matched pattern.
+
+    The `Match` object is returned to represent a successfully matched
+    pattern. Included in the Match are the pattern that was matched, the graph
+    nodes matched, and any args that were used during the matching.
+
+    The args and kwargs are specific to the type of pattern that was matched and
+    provide hints about what was matched.
+    """
+
+    pattern: PatternExpr
+    args: list[Any]
+    kwargs: dict[str, Any]
+    nodes: list[torch.fx.Node]
+    targets: dict[_TargetExpr, torch.fx.node.Target]
+    ctx: MatchContext
+    replacement_graph: torch.fx.GraphModule | None
+
+    def __init__(
+        self,
+        ctx: MatchContext,
+        pattern: PatternExpr,
+        args: Sequence[Any] | None = None,
+        kwargs: dict[str, Any] | None = None,
+    ) -> None:
+        super().__init__()
+        self.pattern = pattern
+        # The input nodes that must be passed in to the result
+        self.args = list(args or [])
+        self.kwargs = kwargs or {}
+        # The nodes matched in this expression
+        self.nodes = []
+        # Mapping CallFunction to the node.target
+        self.targets = {}
+        self.ctx = ctx
+        self.replacement_graph = None
+
+    @property
+    def graph(self) -> torch.fx.Graph:
+        return self.ctx.graph
+
+    def extend(self, other: Match) -> None:
+        if self.kwargs:
+            for key in OrderedSet(self.kwargs.keys()) & OrderedSet(other.kwargs.keys()):
+                if self.kwargs[key] != other.kwargs[key]:
+                    raise FailedMatch("kwarg mismatch: {}", key)
+        self.args.extend(other.args)
+        self.nodes.extend(other.nodes)
+        self.kwargs.update(other.kwargs)
+        self.targets.update(other.targets)
+
+    def bundle(self) -> Match:
+        # Wrap args in an extra list
+        self.args = [tuple(self.args)] if self.args else []
+        return self
+
+    def __repr__(self) -> str:
+        return f"Match(..., {self.args}, {self.kwargs})"
+
+    def erase_nodes(self) -> None:
+        graph = self.graph
+        for n in reversed(self.nodes):
+            if not n._erased and not n.users:
+                graph.erase_node(n)
+
+    def output_nodes(self) -> list[torch.fx.Node | None]:
+        return [
+            (self.ctx.pattern_to_node[p] if p is not None else None)
+            for p in self.ctx.outputs
+        ]
+
+    def output_node(self) -> torch.fx.Node:
+        return next(p for p in self.output_nodes() if p)
+
+    def replace_with_graph(
+        self, replacement_graph: torch.fx.Graph, args: Sequence[Any]
+    ) -> None:
+        ReplacementPatternEntry.replace_with_graph(
+            self, self.ctx.graph, replacement_graph, args
+        )
+
+    def replace_by_example(
+        self,
+        replacement_fn: ReplaceFn,
+        args: Sequence[Any],
+        trace_fn: TraceFn | None = None,
+        run_functional_passes: bool = True,
+    ) -> None:
+        """Replace with a graph generated by tracing the replacement_fn.
+
+        Args:
+            run_functional_passes (bool). If we should run passes that
+                assume functional IR (like DCE, remove_noop_ops), on the
+                replacement graph.
+
+        """
+        from torch._inductor.virtualized import NullHandler, V
+
+        context = (
+            V.fake_mode
+            if (not isinstance(V.fake_mode, NullHandler) or (V.fake_mode is None))
+            else contextlib.nullcontext()
+        )
+
+        def should_propagate_eager_input_vals(nodes: list[torch.fx.Node]) -> bool:
+            if len(nodes) != 1:
+                return False
+            node = nodes[0]
+            if "eager_input_vals" not in node.meta:
+                return False
+            return node.target in OrderedSet(
+                [
+                    torch.ops.higher_order.triton_kernel_wrapper_functional,
+                    torch.ops.higher_order.auto_functionalized,
+                    torch.ops.higher_order.auto_functionalized_v2,
+                ]
+            )
+
+        # pyrefly: ignore [bad-context-manager]
+        with context:
+            if trace_fn is None:
+                trace_fn = functools.partial(
+                    fwd_only, run_functional_passes=run_functional_passes
+                )
+
+            if should_propagate_eager_input_vals(self.nodes):
+                # Our strategy is:
+                # 1) trace out the graph with eager_input_vals (which have accurate eager-mode metadata)
+                # 2) trace out the graph with vals (which have the accurate Inductor metadata)
+                # 3) Propagate the eager_input_vals from the first graph to the second.
+                # 4) Use the second graph as the replacement graph.
+
+                # Construct a map of node -> FakeTensor val in eager_input_vals
+                node_to_val = {}
+
+                fake_args, fake_kwargs = self.nodes[0].meta["eager_input_vals"]
+                fake_kwargs = {**fake_kwargs}
+                match_args, match_kwargs = tuple(self.args), self.kwargs
+
+                def record(node: torch.fx.Node, val: Any) -> None:
+                    if isinstance(node, torch.fx.Node):
+                        node_to_val[node] = val
+
+                torch.utils._pytree.tree_map(
+                    record, (match_args, match_kwargs), (fake_args, fake_kwargs)
+                )
+                # map args to their FakeTensor val in eager_input_vals
+                example_vals = torch.fx.map_arg(args, lambda arg: node_to_val[arg])
+
+                # first graph
+                graph_with_eager_vals = trace_fn(replacement_fn, example_vals)
+
+                # second graph
+                example_vals = torch.fx.map_arg(args, lambda arg: arg.meta["val"])
+                # pyrefly: ignore [bad-argument-type]
+                replacement = trace_fn(graph_with_eager_vals, example_vals)
+
+                # propagate metadata from first graph to second
+                # NB: This assertion might not be true in general, but it is true for
+                # the two use cases we have
+                # (triton_kernel_wrapper_functional, auto_functionalized)
+                assert len(graph_with_eager_vals.graph.nodes) == len(
+                    replacement.graph.nodes
+                )
+                for old_node, new_node in zip(
+                    graph_with_eager_vals.graph.nodes, replacement.graph.nodes
+                ):
+                    if "eager_input_vals" in old_node.meta:
+                        new_node.meta["eager_input_vals"] = old_node.meta[
+                            "eager_input_vals"
+                        ]
+
+            else:
+                example_vals = torch.fx.map_arg(
+                    args,
+                    lambda arg: arg.meta["val"]
+                    if "val" in arg.meta
+                    else arg.meta["example_value"],
+                )
+                fake_mode = torch._dynamo.utils.detect_fake_mode(example_vals)
+                if fake_mode is not None:
+
+                    def _convert_to_fake_mode(it):
+                        if isinstance(it, FakeTensor) and fake_mode is not None:
+                            return fake_mode.from_tensor(it)
+                        return it
+
+                    example_vals = torch.fx.node.map_aggregate(
+                        example_vals, _convert_to_fake_mode
+                    )
+                replacement = trace_fn(replacement_fn, example_vals)
+            if len(self.nodes) == 1:
+                for n in replacement.graph.nodes:
+                    _transfer_meta(
+                        new_meta=n.meta,
+                        old_node=self.nodes[0],
+                        pass_name="replace_by_example",
+                    )
+
+            ReplacementPatternEntry.replace_with_graph(
+                self,
+                self.ctx.graph,
+                replacement,
+                args,
+            )
+
+
+class FailedMatch(RuntimeError):
+    """
+    Represents a unsuccessful match.
+
+    The `FailedMatch` object is returned to represent a failure to match a
+    pattern.
+    """
+
+    format_string: str
+
+    def __init__(self, format_string: str, *args: Any, **kwargs: Any) -> None:
+        self.format_string = format_string
+        # We want to construct error messages lazily instead of eagerly, as
+        # constructing them eagerly can significantly worsen compile times.
+        if len(format_string) > 200:
+            raise RuntimeError(
+                f"Format string too long - use lazy construction of strings instead. Format string is\n {format_string}"
+            )
+        self.args = args
+        self.kwargs = kwargs
+
+    def __str__(self) -> str:
+        return self.format_string.format(*self.args, **self.kwargs)
+
+    def __bool__(self) -> bool:
+        return False
+
+
+MatchResult = Match | FailedMatch
+
+
+def is_match(m: MatchResult) -> TypeIs[Match]:
+    """
+    TypeIs cannot act on `self`. Thus this function exists to let mypy
+    recognize FailedMatch.__bool__ as a TypeIs.
+    """
+    return bool(m)
+
+
+class MatchContext:
+    """
+    Internal state needed while running PatternExpr._match().
+    """
+
+    outputs: list[PatternExpr | None]
+    pattern_to_node: dict[PatternExpr, torch.fx.Node | None]
+    graph: torch.fx.Graph
+    exclusive_node_set: list[NodeOrConstant]
+
+    def __init__(
+        self,
+        outputs: list[PatternExpr | None],
+        pattern_to_node: dict[PatternExpr, torch.fx.Node] | None = None,
+        *,
+        graph: torch.fx.Graph,
+    ) -> None:
+        self.outputs = outputs
+        self.pattern_to_node = {} if pattern_to_node is None else dict(pattern_to_node)
+        self.graph = graph
+        self.exclusive_node_set = []
+
+    def match(self, pattern: PatternExpr, node: NodeOrConstant) -> MatchResult:
+        """wrapper to check reused nodes in patterns"""
+        if pattern in self.pattern_to_node:
+            if self.pattern_to_node[pattern] == node:
+                return Match(self, pattern)  # already checked this node
+            else:
+                return FailedMatch("repeated pattern differs")
+        m = pattern._match(node, self)
+        assert pattern not in self.pattern_to_node
+        self.pattern_to_node[pattern] = node if m else None
+        return m
+
+    def filter_multi_user_patterns(self) -> dict[PatternExpr, torch.fx.Node]:
+        return {
+            pattern: node
+            for pattern, node in self.pattern_to_node.items()
+            if pattern.has_multiple_users() and node is not None
+        }
+
+
+class PatternExpr(ABC):
+    """
+    Base class for types of patterns.
+    """
+
+    @abstractmethod
+    def _match(self, node: torch.fx.Node, ctx: MatchContext) -> MatchResult: ...
+
+    def match(self, node: torch.fx.Node) -> MatchResult:
+        try:
+            return MatchContext([self], graph=node.graph).match(self, node)
+        except FailedMatch as e:
+            return e
+
+    def has_multiple_users(self) -> bool:
+        return False
+
+    def __repr__(self) -> str:
+        return self.__class__.__name__ + "()"
+
+    def find_anchor_nodes(
+        self, ctx: MatchContext, searched: OrderedSet[torch.fx.Node]
+    ) -> Generator[torch.fx.Node | None, None, None]:
+        if self in ctx.pattern_to_node:
+            yield ctx.pattern_to_node[self]
+
+    def pattern_eq(self, other: Any) -> bool:
+        """
+        Compare two `PatternExpr`s and return true if they are the
+        same. Note this is NOT matching a pattern - it is comparing the pattern
+        structures (for debugging).
+        """
+        return isinstance(other, self.__class__)
+
+
+class Arg(PatternExpr):
+    """
+    Capture an arg which will become an input to the handler.  Args are
+    passed in depth first order.
+    """
+
+    def _match(self, node: NodeOrConstant, ctx: MatchContext) -> MatchResult:
+        return Match(ctx, self, args=[node])  # matches anything
+
+
+class Ignored(PatternExpr):
+    """
+    Match an arg, but don't pass it to handler
+    """
+
+    def _match(self, node: NodeOrConstant, ctx: MatchContext) -> MatchResult:
+        return Match(ctx, self)  # matches anything
+
+    def __repr__(self) -> str:
+        return "*"
+
+    def pretty_print(self, pp: PatternPrettyPrinter) -> str:
+        return "Ignored()"
+
+
+class KeywordArg(PatternExpr):
+    """
+    Capture a kwarg which will become an input to the handler.
+    """
+
+    def __init__(self, name: str) -> None:
+        super().__init__()
+        self.name = name
+
+    def __repr__(self) -> str:
+        return f"KeywordArg({self.name!r})"
+
+    def _match(self, node: NodeOrConstant, ctx: MatchContext) -> MatchResult:
+        return Match(ctx, self, kwargs={self.name: node})  # matches anything
+
+    def pattern_eq(self, other: Any) -> bool:
+        other = typing.cast(Self, other)  # super makes sure this is true
+        return super().pattern_eq(other) and self.name == other.name
+
+
+class ExclusiveKeywordArg(PatternExpr):
+    """
+    Capture a kwarg which will become an input to the handler.
+    """
+
+    name: str
+
+    def __init__(self, name: str) -> None:
+        super().__init__()
+        self.name = name
+
+    def __repr__(self) -> str:
+        return f"ExclusiveKeywordArg({self.name!r})"
+
+    def _match(self, node: NodeOrConstant, ctx: MatchContext) -> MatchResult:
+        if node in ctx.exclusive_node_set:
+            return FailedMatch("exclusive arg appears twice")
+
+        ctx.exclusive_node_set.append(node)
+        return Match(ctx, self, kwargs={self.name: node})  # matches anything
+
+    def pattern_eq(self, other: Any) -> bool:
+        other = typing.cast(Self, other)  # super makes sure this is true
+        return super().pattern_eq(other) and self.name == other.name
+
+
+class _TargetExpr(PatternExpr):
+    """
+    Base class for filtering match by node.target
+    """
+
+    fns: list[FnsType]
+    fns_set: OrderedSet[FnsType]
+
+    def __init__(
+        self, fns: FnsType | Sequence[FnsType], users: Multiple | int = 1
+    ) -> None:
+        super().__init__()
+        fns = [fns] if callable(fns) or isinstance(fns, str) else list(fns)
+        for fn in fns:
+            if isinstance(fn, torch._ops.OpOverloadPacket):
+                fns.extend(getattr(fn, overload) for overload in fn.overloads())  # noqa: B909
+
+        self.fns = fns
+        self.fns_set = OrderedSet(fns)
+        self.users = users
+
+    @property
+    @abstractmethod
+    def op(self) -> str: ...
+
+    def fns_repr(self) -> str:
+        first_repr = self.fns[0]
+        if not isinstance(first_repr, str):
+            first_repr = first_repr.__name__
+
+        if len(self.fns) > 1:
+            return f"[{first_repr}, ...]"
+        elif self.fns[0] is getattr(torch, first_repr, None):
+            return f"torch.{first_repr}"
+        elif self.fns[0] is getattr(operator, first_repr, None):
+            return f"operator.{first_repr}"
+        elif isinstance(self.fns[0], torch._ops.OpOverload):
+            return str(self.fns[0])
+        else:
+            return first_repr
+
+    def __repr__(self) -> str:
+        if self.users is MULTIPLE:
+            comma_users = ", MULTIPLE"
+        elif self.users != 1:
+            comma_users = f", {self.users})"
+        else:
+            comma_users = ""
+        return f"{self.__class__.__name__}({self.fns_repr()}{comma_users})"
+
+    def has_multiple_users(self) -> bool:
+        return isinstance(self.users, Multiple) or self.users > 1
+
+    def find_anchor_nodes(
+        self, ctx: MatchContext, searched: OrderedSet[torch.fx.Node]
+    ) -> Generator[torch.fx.Node | None, None, None]:
+        raise NotImplementedError
+
+    def _match_fns(self, node: torch.fx.Node) -> bool:
+        return (
+            isinstance(node, torch.fx.Node)
+            and node.op == self.op
+            and extract_target(node) in self.fns_set
+        )
+
+    def _match_users(self, node: torch.fx.Node, ctx: MatchContext) -> bool:
+        return (
+            self in ctx.outputs
+            or self.users is MULTIPLE
+            or len(node.users) == self.users
+        )
+
+    def pattern_eq(self, other: Any) -> bool:
+        other = typing.cast(Self, other)  # super makes sure this is true
+        return (
+            super().pattern_eq(other)
+            and self.op == other.op
+            and self.fns == other.fns
+            and self.users == other.users
+        )
+
+
+_SimpleSpec = tuple[Any, ...]
+
+
+class _TargetArgsExpr(_TargetExpr):
+    """
+    Base class for filtering match by node.{target,args,kwargs}
+    """
+
+    def __init__(
+        self,
+        fns: torch.fx.node.Target | str | Sequence[Any],
+        *args: Any,
+        _users: int | Multiple = 1,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(fns, _users)
+        self.args = tuple(args)
+        self.kwargs = dict(kwargs)
+        if any(
+            isinstance(x, (dict, list, tuple))
+            for x in itertools.chain(args, kwargs.values())
+        ):
+            self.flatten = self.pytree_flatten
+        else:
+            self.flatten = self.simple_flatten
+        self.flat_args_kwargs = self.flatten(self.args, self.kwargs)
+
+    @staticmethod
+    def simple_flatten(
+        args: Sequence[Any], kwargs: Mapping[Any, Any]
+    ) -> tuple[Sequence[Any], _SimpleSpec | pytree.TreeSpec]:
+        values = (*args, *kwargs.values())
+        spec = (len(args), *kwargs.keys())
+        return values, spec
+
+    @staticmethod
+    def pytree_flatten(
+        args: Sequence[Any], kwargs: Mapping[Any, Any]
+    ) -> tuple[Sequence[Any], _SimpleSpec | pytree.TreeSpec]:
+        type_mapping: dict[type, type] = {
+            immutable_list: tuple,
+            list: tuple,
+            immutable_dict: dict,
+        }
+
+        def convert_type(x: Any) -> Any:
+            cls = type(x)
+            convert_fn = type_mapping.get(cls)
+            if convert_fn is not None:
+                return pytree.tree_map(
+                    convert_type,
+                    convert_fn(x),
+                    is_leaf=lambda x: type(x) in type_mapping,
+                )
+            return x
+
+        normalized_args_tree = pytree.tree_map(
+            convert_type,
+            (args, kwargs),
+            is_leaf=lambda x: type(x) in type_mapping,
+        )
+        flat, spec = pytree.tree_flatten(normalized_args_tree)
+        return flat, spec
+
+    def __repr__(self) -> str:
+        args = [
+            self.fns_repr(),
+            *map(repr, self.args),
+            *[f"{k}={v}" for k, v in self.kwargs.items()],
+        ]
+        if self.users is MULTIPLE:
+            args.append("_users=MULTIPLE")
+        elif self.users != 1:
+            args.append(f"_users={self.users}")
+        return f"{self.__class__.__name__}({', '.join(args)})"
+
+    def pretty_print(self, pp: PatternPrettyPrinter) -> str:
+        args = [
+            self.fns_repr(),
+            *(pp.pretty_print(x) for x in self.args),
+            *[f"{k}={pp.pretty_print(v)}" for k, v in self.kwargs.items()],
+        ]
+        if self.users is MULTIPLE:
+            args.append("_users=MULTIPLE")
+        elif self.users != 1:
+            args.append(f"_users={self.users}")
+
+        joiner_str = ", "
+        return f"{self.__class__.__name__}({joiner_str.join(args)})"
+
+    def _match(self, node: torch.fx.Node, ctx: MatchContext) -> MatchResult:
+        if not self._match_fns(node) or len(node.args) != len(self.args):
+            return FailedMatch("function_mismatch: node={}, pattern={}", node, self)
+
+        if not self._match_users(node, ctx):
+            return FailedMatch("multiple_users {}", self)
+
+        _args = node.args
+        _kwargs = node.kwargs
+        if len(_kwargs) < len(self.kwargs):
+            from torch.fx.operator_schemas import normalize_function
+
+            assert callable(node.target)
+            normalized_args_and_kwargs = normalize_function(
+                node.target, node.args, node.kwargs
+            )
+
+            if normalized_args_and_kwargs is None:
+                return FailedMatch("function_mismatch: node={}, pattern={}", node, self)
+            else:
+                _args, _kwargs = normalized_args_and_kwargs
+                if len(_args) == len(self.args) and len(_kwargs) >= len(self.kwargs):
+                    _kwargs = {i: _kwargs[i] for i in _kwargs if i in self.kwargs}
+                else:
+                    return FailedMatch(
+                        "function_mismatch: node={}, pattern={}", node, self
+                    )
+        else:
+            _kwargs = {i: _kwargs[i] for i in _kwargs if i in self.kwargs}
+
+        node_items, node_spec = self.flatten(_args, _kwargs)
+        self_items, self_spec = self.flat_args_kwargs
+        if node_spec != self_spec:
+            return FailedMatch("args_structure {} {}", node_spec, self_spec)
+        assert len(node_items) == len(self_items)
+
+        m = Match(ctx, self)
+        for pattern, child_node in zip(self_items, node_items):
+            if isinstance(pattern, PatternExpr):
+                child_match = ctx.match(pattern, child_node)
+                if not is_match(child_match):
+                    return child_match
+                m.extend(child_match)
+            elif isinstance(child_node, torch.fx.Node) or child_node != pattern:
+                return FailedMatch(
+                    "constant_args: {} {!r}!={pattern!r}",
+                    node,
+                    child_node,
+                    pattern=pattern,
+                )
+        m.nodes.append(node)
+        m.targets[self] = node.target
+        return m
+
+    def find_anchor_nodes(
+        self, ctx: MatchContext, searched: OrderedSet[torch.fx.Node]
+    ) -> Generator[torch.fx.Node | None, None, None]:
+        """
+        This is used when we are matching a pattern with multiple outputs.
+        There is a partial match (stored in ctx) and we want to walk
+        this pattern to find a connection to an already-matched node.
+
+        Yields candidate nodes that `self._match` might like.
+        """
+        if self in ctx.pattern_to_node:
+            yield ctx.pattern_to_node[self]
+            return
+
+        for pattern in self.flat_args_kwargs[0]:
+            if isinstance(pattern, PatternExpr):
+                for other_node in pattern.find_anchor_nodes(ctx, searched):
+                    if not isinstance(other_node, torch.fx.Node):
+                        continue
+                    for node in other_node.users:
+                        if node not in searched:
+                            if self._match_fns(node):
+                                yield node
+                                searched.add(node)
+
+    def pattern_eq(self, other: Any) -> bool:
+        other = typing.cast(Self, other)  # super makes sure this is true
+        return (
+            super().pattern_eq(other)
+            and self.flat_args_kwargs[1] == other.flat_args_kwargs[1]
+            and all(
+                a.pattern_eq(b) if isinstance(a, PatternExpr) else a == b
+                for a, b in zip(self.flat_args_kwargs[0], other.flat_args_kwargs[0])
+            )
+        )
+
+
+class CallFunction(_TargetArgsExpr):
+    """
+    Matches a call_function node in the FX graphs: `fns[i](*args, **kwargs)`
+    """
+
+    op = "call_function"
+
+
+class CallMethod(_TargetArgsExpr):
+    """
+    Matches a call_method node in the FX graphs: `fns[i].method(*args, **kwargs)`
+    """
+
+    op = "call_method"
+
+
+class CallModule(_TargetArgsExpr):
+    """
+    Matches a call_module node in the FX graphs: `module(*args, **kwargs)`
+    """
+
+    op = "call_module"
+
+
+class _TargetExprVarArgs(_TargetExpr):
+    """
+    Matches a call_function node with any arguments which are passed into the pattern
+    """
+
+    def _match(self, node: torch.fx.Node, ctx: MatchContext) -> MatchResult:
+        if not self._match_fns(node):
+            return FailedMatch("function_mismatch")
+
+        if not self._match_users(node, ctx):
+            return FailedMatch("multiple_users")
+
+        m = Match(ctx, self)
+        m.nodes.append(node)
+        m.targets[self] = node.target
+        m.args.extend(node.args)
+        m.kwargs.update(node.kwargs)
+        return m
+
+
+class CallFunctionVarArgs(_TargetExprVarArgs):
+    op = "call_function"
+
+
+class CallMethodVarArgs(_TargetExprVarArgs):
+    op = "call_method"
+
+
+class CallModuleVarArgs(_TargetExprVarArgs):
+    op = "call_module"
+
+
+class ListOf(PatternExpr):
+    """
+    Matches a repeated pattern
+    """
+
+    def __init__(self, pattern: PatternExpr, partial: bool = False) -> None:
+        super().__init__()
+        assert isinstance(pattern, PatternExpr)
+        self.pattern = pattern
+        self.partial = partial
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({self.pattern})"
+
+    def _match(self, node: list[torch.fx.Node], ctx: MatchContext) -> MatchResult:  # type: ignore[override]
+        if not isinstance(node, (list, tuple)) or len(node) == 0:
+            return FailedMatch("non_list")
+        m = Match(ctx, self)
+        # Propagating patterns with multiple users will ensure we don't revisit
+        # the same nodes
+        pattern_to_node = ctx.filter_multi_user_patterns()
+        matched = False
+        for i, child_node in enumerate(node):
+            child_ctx = MatchContext(
+                ctx.outputs, pattern_to_node, graph=child_node.graph
+            )
+            child_match = child_ctx.match(self.pattern, child_node)
+            pattern_to_node = child_ctx.filter_multi_user_patterns()
+            if not is_match(child_match):
+                if not self.partial:
+                    return FailedMatch("list[{}]: {}", i, child_match)
+                continue
+            matched = True
+            m.extend(child_match.bundle())
+        if not matched:
+            return FailedMatch("list: no_match")
+        return m.bundle()
+
+    def pattern_eq(self, other: Any) -> bool:
+        other = typing.cast(Self, other)  # super makes sure this is true
+        return (
+            super().pattern_eq(other)
+            and self.pattern.pattern_eq(other.pattern)
+            and self.partial == other.partial
+        )
+
+
+class MultiOutputPattern(PatternExpr):
+    outputs: list[PatternExpr | None]
+
+    def __init__(self, outputs: Sequence[PatternExpr | None]) -> None:
+        super().__init__()
+        assert isinstance(outputs[0], _TargetExpr)
+        assert all(x is None or isinstance(x, PatternExpr) for x in outputs), outputs
+        self.outputs = list(outputs)
+        self.op = outputs[0].op
+
+    @property
+    def fns(self) -> Callable[..., Any] | str | Sequence[Any]:
+        # This cast is checked above in __init__()
+        output = typing.cast(_TargetExpr, self.outputs[0])
+        return output.fns
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({self.outputs})"
+
+    def pretty_print(self, pp: PatternPrettyPrinter) -> str:
+        args = [pp.pretty_print(x) for x in self.outputs]
+        joiner_str = f",\n{'  '}"
+        str_out = f"{self.__class__.__name__}([{joiner_str.join(args)}"
+        str_out = f"{str_out}\n])"
+        return str_out
+
+    def _match(self, node: torch.fx.Node, ctx: MatchContext) -> MatchResult:
+        output = typing.cast(_TargetExpr, self.outputs[0])
+        m = ctx.match(output, node)
+        if not is_match(m):
+            return m
+
+        for pattern in self.outputs[1:]:
+            if pattern is None:
+                continue
+            child_match = self._match_from_anchors(pattern, ctx)
+            if not is_match(child_match):
+                return child_match
+            m.extend(child_match)
+
+        return m
+
+    def _match_from_anchors(
+        self, pattern: PatternExpr, ctx: MatchContext
+    ) -> MatchResult:
+        prior = dict(ctx.pattern_to_node)
+        m: MatchResult = FailedMatch("no anchor found")
+        for node in pattern.find_anchor_nodes(ctx, OrderedSet()):
+            m = ctx.match(pattern, node)
+            if is_match(m):
+                return m
+            # revert any partial matches
+            ctx.pattern_to_node = dict(prior)
+        return m
+
+    def match(self, node: torch.fx.Node) -> MatchResult:
+        try:
+            return MatchContext(self.outputs, graph=node.graph).match(self, node)
+        except FailedMatch as e:
+            return e
+
+    def pattern_eq(self, other: Any) -> bool:
+        other = typing.cast(Self, other)  # super makes sure this is true
+        return (
+            super().pattern_eq(other)
+            and len(self.outputs) == len(other.outputs)
+            and all(
+                a.pattern_eq(b) if isinstance(a, PatternExpr) else a == b
+                for a, b in zip(self.outputs, other.outputs)
+            )
+        )
+
+
+class RepeatedExpr(PatternExpr):
+    """
+    Checks for a repeated pattern. Useful for repeated operations after a node such as `split` or `unbind`
+    """
+
+    def __init__(self, inner_pattern: _TargetExpr) -> None:
+        super().__init__()
+        self.inner_pattern = inner_pattern
+        self.op = inner_pattern.op
+
+    @property
+    def fns(self) -> Sequence[FnsType]:
+        return self.inner_pattern.fns
+
+    def _match(self, node: torch.fx.Node, ctx: MatchContext) -> MatchResult:
+        m = ctx.match(self.inner_pattern, node)
+        if not is_match(m):
+            return m
+        ctx.pattern_to_node.pop(
+            self.inner_pattern,
+        )
+        # Check all anchor nodes match the pattern
+        for anchor_node in self.inner_pattern.find_anchor_nodes(
+            ctx, OrderedSet([node])
+        ):
+            anchor_m = MatchContext([self], graph=node.graph).match(
+                self.inner_pattern, anchor_node
+            )
+            if not is_match(anchor_m):
+                return anchor_m
+            m.extend(anchor_m)
+        return m
+
+    def pattern_eq(self, other: Any) -> bool:
+        other = typing.cast(Self, other)  # super makes sure this is true
+        return super().pattern_eq(other) and self.inner_pattern.pattern_eq(
+            other.inner_pattern
+        )
+
+
+class PatternPrettyPrinter:
+    """
+    Serializes Patterns to executable python.
+    XXX: currently only used and tested for fuse attention patterns. May not cover
+    all patterns.
+    """
+
+    def __init__(self) -> None:
+        self.namespace = torch.fx.graph._Namespace()
+        self.memoized_objs_names: dict[PatternExpr, str] = {}
+        self.memoized_objs_pp: dict[PatternExpr, str] = {}
+
+    @staticmethod
+    @functools.cache
+    def run(obj: PatternExpr, output_name: str = "output") -> str:
+        """
+        Serializes obj to python code with obj written out to `output_name`
+        """
+
+        pp = PatternPrettyPrinter()
+        assert hasattr(obj, "pretty_print")
+        out_str = obj.pretty_print(pp=pp)
+
+        output = [
+            f"{pp.memoized_objs_names[key]} = {pp.memoized_objs_pp[key]}"
+            for key in pp.memoized_objs_names
+        ]
+
+        output.append(f"{output_name} = {out_str}")
+
+        return "\n".join(output)
+
+    def pretty_print(self, obj: Any) -> str:
+        if isinstance(obj, _TargetArgsExpr):
+            if memoized_name := self.memoized_objs_names.get(obj):
+                return memoized_name
+            else:
+                return self.memoize(obj)
+        if hasattr(obj, "pretty_print"):
+            return obj.pretty_print(self)
+
+        return repr(obj)
+
+    def memoize(self, obj: _TargetArgsExpr) -> str:
+        obj_str = obj.pretty_print(self)
+        obj_name = obj.fns_repr()
+        for prefix in ("aten.", "torch.", "prims."):
+            obj_name = obj_name.replace(prefix, "")
+
+        tmp_name = self.namespace.create_name(obj_name, None)
+        self.memoized_objs_names[obj] = tmp_name
+        self.memoized_objs_pp[obj] = obj_str
+        return tmp_name
+
+
+class _PassDictsType(Protocol):
+    def __getitem__(
+        self, k: tuple[str, torch.fx.node.Target]
+    ) -> list[PatternEntry]: ...
+
+
+@dataclasses.dataclass
+class PatternEntry:
+    pattern: PatternExpr
+    extra_check: Callable[[Match], bool]
+
+    def apply(self, match: Match, graph: torch.fx.Graph, node: torch.fx.Node) -> None:
+        raise NotImplementedError
+
+    def register(
+        self,
+        pass_dicts: _PassDictsType | Sequence[_PassDictsType],
+        target: torch.fx.node.Target | None = None,
+        prepend: bool = False,
+    ) -> None:
+        if target is None:
+            assert hasattr(self.pattern, "fns")
+            for fn in self.pattern.fns:
+                self.register(pass_dicts, fn, prepend=prepend)
+        elif isinstance(pass_dicts, (dict, PatternMatcherPass)):
+            assert hasattr(self.pattern, "op")
+            if prepend:
+                pass_dicts[(self.pattern.op, target)].insert(0, self)
+            else:
+                pass_dicts[(self.pattern.op, target)].append(self)
+        else:
+            pass_dicts = typing.cast(Sequence[_PassDictsType], pass_dicts)
+            for x in pass_dicts:
+                self.register(x, target, prepend=prepend)
+
+
+@dataclasses.dataclass
+class LoweringPatternEntry(PatternEntry):
+    handler: Callable[..., Any]
+
+    def apply(self, match: Match, graph: torch.fx.Graph, node: torch.fx.Node) -> None:
+        handler = functools.wraps(self.handler)(functools.partial(self.handler, match))
+        with graph.inserting_before(node):
+            replacement = graph.call_function(handler, tuple(match.args), match.kwargs)
+            _transfer_meta(replacement.meta, node)
+            node.replace_all_uses_with(replacement)
+        assert match.nodes[-1] is node
+        match.erase_nodes()
+
+
+@dataclasses.dataclass
+class GraphPatternEntry(PatternEntry):
+    """
+    A pattern that runs a function on the FX graph
+    """
+
+    handler: Callable[..., Any]
+
+    def apply(self, match: Match, graph: torch.fx.Graph, node: torch.fx.Node) -> None:
+        with graph.inserting_before(node):
+            self.handler(match, *match.args, **match.kwargs)
+
+
+@dataclasses.dataclass
+class ReplacementPatternEntry(PatternEntry):
+    """
+    The replacement pattern for the graph
+    """
+
+    normalize_args: Callable[..., list[Any]]
+    pattern_name: str | None = None  # Unique identifier for per-pattern telemetry
+
+    @staticmethod
+    def replace_with_graph(
+        match: Match,
+        graph: torch.fx.Graph,
+        replacement_graph: torch.fx.Graph | torch.fx.GraphModule,
+        args: Sequence[torch.fx.Node],
+        pass_name: str | None = None,
+    ) -> None:
+        """
+        Inserts the replacement graph into the toplevel graph at the match
+        """
+
+        added_replacement_nodes: list[torch.fx.Node] = []
+
+        class Replacer(torch.fx.Interpreter):
+            call_method = None  # type: ignore[assignment]
+            call_module = None  # type: ignore[assignment]
+            get_attr = None  # type: ignore[assignment]
+
+            def run_node(self, node: torch.fx.Node) -> Any:
+                if node.op in ("placeholder", "output"):
+                    return super().run_node(node)
+                target = node.target
+                args, kwargs = self.fetch_args_kwargs_from_env(node)
+                if node.op == "call_function":
+                    assert callable(target)
+                    result = graph.call_function(target, args, kwargs)
+                    added_replacement_nodes.append(result)
+                    _transfer_meta(
+                        new_meta=result.meta,
+                        old_node=node,
+                        pass_name=pass_name or "",
+                    )
+                    # This function copy-pastes the replacement graph into
+                    # the graph. If the replacement graph had any eager_input_vals,
+                    # we propagate those over (val/tensor_meta are handled by
+                    # _transfer_meta above).
+                    if "eager_input_vals" in node.meta:
+                        result.meta["eager_input_vals"] = node.meta["eager_input_vals"]
+                    return result
+                if node.op == "get_attr":
+                    # If the replacement graph contains a HOP, the subgraphs of the HOP are "get_attr" nodes.
+                    # We need to fetch the subgraph of the HOP then register the subgraph to the replaced graph's root.
+                    from torch._higher_order_ops.utils import (
+                        unique_graph_name_with_root,
+                    )
+
+                    sub_gm = super().get_attr(target, args, kwargs)
+                    if not isinstance(sub_gm, torch.fx.GraphModule):
+                        raise NotImplementedError(
+                            f"NYI: replacement_graph.{target} is not a graph module. Got {sub_gm}."
+                        )
+                    assert graph.owning_module is not None
+                    graph_name = None
+                    for n, mod in graph.owning_module.named_modules():
+                        if sub_gm is mod:
+                            graph_name = n
+                            break
+                    if graph_name is None:
+                        assert isinstance(target, str)
+                        _, graph_name = unique_graph_name_with_root(
+                            # pyrefly: ignore [unbound-name]
+                            graph.owning_module,
+                            target,
+                        )
+                        # pyrefly: ignore [unbound-name]
+                        graph.owning_module.register_module(graph_name, sub_gm)
+                    # pyrefly: ignore [unbound-name]
+                    getattr_node = graph.get_attr(graph_name)
+                    added_replacement_nodes.append(getattr_node)
+                    return getattr_node
+
+                raise NotImplementedError(f"unhandled {node}")
+
+        output_nodes = match.output_nodes()
+
+        if len(output_nodes) == 1:
+            last_node = output_nodes[0]
+        else:
+            assert output_nodes[0]
+            nodes = list(output_nodes[0].graph.nodes)
+            indices = [
+                (nodes.index(n), n)
+                for n in output_nodes
+                if isinstance(n, torch.fx.Node)
+            ]
+            last_node = min(indices, key=operator.itemgetter(0))[1]
+
+        def percolate_tags(
+            node: torch.fx.Node,
+            tag_name: str,
+            tag_value: str,
+            input_stops: OrderedSet[torch.fx.Node],
+        ) -> None:
+            queue = [node]
+            visited = OrderedSet[torch.fx.Node]()
+
+            while queue:
+                arg = queue.pop()
+                if (
+                    arg not in visited
+                    and arg not in input_stops
+                    and hasattr(arg, "meta")
+                ):
+                    visited.add(arg)
+                    arg.meta[tag_name] = tag_value
+                    queue.extend(arg.all_input_nodes)
+
+        with graph.inserting_before(last_node):
+            assert isinstance(replacement_graph, torch.fx.GraphModule)
+            replacement = Replacer(replacement_graph).run(*args)
+            if isinstance(replacement, torch.fx.Node):
+                replacement = [replacement]
+
+            def maybe_getitem(node: torch.fx.Node) -> Any:
+                if node.op != "call_function":
+                    return None
+                if node.target != operator.getitem:
+                    return None
+                assert len(node.args) == 2
+                return node.args[1]
+
+            def replace(
+                old: torch.fx.Node | None,
+                new: torch.fx.Node | Sequence[torch.fx.Node] | None,
+            ) -> None:
+                def filter_nodes_in_newly_added_nodes(node: torch.fx.Node) -> bool:
+                    # Do not replace the use of a node if it is being used by
+                    # nodes in the replaced graph
+                    return node not in added_replacement_nodes
+
+                if old is None:
+                    assert new is None
+                    return
+                assert isinstance(old, torch.fx.Node)
+                if new is None:
+                    old.replace_all_uses_with(
+                        None,  # type: ignore[arg-type]
+                        delete_user_cb=filter_nodes_in_newly_added_nodes,
+                    )
+                    if len(old.users) == 0:
+                        graph.erase_node(old)
+                    return
+                if isinstance(new, torch.fx.Node):
+                    _transfer_meta(new.meta, old, pass_name=pass_name or "")
+
+                    # Preserve the recompute tags in the replacement graph. We
+                    # look at the recompute tags of the original output node to
+                    # propagate the tag from the output all the way to the input
+                    # args (named as args in the replace_with_graph).
+                    # Note that this is best effort. Since patterns are from
+                    # many to many, there is no easy way to correctly map the
+                    # recomputable tags. It is possible in some scenarios that we
+                    # incorrectly tag some nodes as recomputables.
+                    for tag_name in ["recompute", "ac_graph_id"]:
+                        if tag_name in old.meta:
+                            percolate_tags(
+                                new, tag_name, old.meta[tag_name], OrderedSet(args)
+                            )
+
+                    old.replace_all_uses_with(
+                        new, delete_user_cb=filter_nodes_in_newly_added_nodes
+                    )
+                    if len(old.users) == 0:
+                        graph.erase_node(old)
+                    return
+
+                # `new` is not a node: it's a list of nodes.
+                #
+                # This happens when we want to replace a node that has a single
+                # packed return with multiple unpacked returns. We need to do
+                # some graph surgery here.
+                #
+                # Example:
+                #   def original_graph(x):
+                #      a = op(x)
+                #      b = a[0]
+                #      c = a[1]
+                #      ...
+                #
+                # Assume that we want to replace op(x) with the graph
+                #   def new_op(x):
+                #      w = x + 1
+                #      z = x + 2
+                #      return (w, z)
+                #
+                # We need to replace `op` with the contents of `new_op`,
+                # and then rewrite a[0] to be w and a[1] to be z, as so:
+                #   def new_graph(x):
+                #     w = x + 1
+                #     z = x + 2
+                #     b = w
+                #     c = z
+                #     ...
+                old_uses = list(old.users.keys())
+                for user in old_uses:
+                    idx = maybe_getitem(user)
+                    if idx is None:
+                        # Output is used directly
+                        # pyrefly: ignore [bad-argument-type]
+                        old.replace_all_uses_with(new)
+                    else:
+                        replace(user, new[idx])
+                graph.erase_node(old)
+
+            if len(output_nodes) == len(replacement):
+                for old, new in zip(output_nodes, replacement):
+                    replace(old, new)
+            else:
+                assert len(output_nodes) == 1
+                replace(output_nodes[0], replacement)
+
+        match.erase_nodes()
+
+        # Remove dead replacement nodes so they don't inflate user counts
+        # in later lowering heuristics (e.g. should_realize_on_reuse).
+        for node in reversed(added_replacement_nodes):
+            if (
+                not node.users
+                and not node.is_impure()
+                and not isinstance(node.target, torch._ops.HigherOrderOperator)
+            ):
+                graph.erase_node(node)
+
+    def apply(self, match: Match, graph: torch.fx.Graph, node: torch.fx.Node) -> None:
+        assert match.replacement_graph is not None
+        self.replace_with_graph(
+            match,
+            graph,
+            match.replacement_graph,
+            self.normalize_args(*match.args, **match.kwargs),
+            pass_name=self.pattern_name or "replace_with_graph",
+        )
+
+
+def _return_true(match: Match) -> bool:
+    return True
+
+
+def log_trace_failure(search_fn: Callable[..., Any], e: RuntimeError) -> None:
+    log.info(
+        "Replacement pattern %s failed to apply due to shape mismatch: %s",
+        search_fn.__name__,
+        e,
+    )
+
+
+def check_and_add_duplicate_pattern(
+    pattern: PatternExpr,
+    graph: torch.fx.Graph | None,
+    seen_patterns: dict[str, list[str | None]],
+    skip_duplicates: bool = False,
+) -> bool:
+    """
+    Check if a pattern is a duplicate. Because we ignore certain types in searching, but not
+    in matching, use the graph to distinguish equivalent search patterns.
+
+    Returns True if a duplicate is found and `skip_duplicates=True` is passed in. Errors if
+    `skip_duplicates` is False and a duplicate is found.
+    """
+
+    pattern_repr = PatternPrettyPrinter.run(pattern)
+    equiv_pattern_reprs = seen_patterns.get(pattern_repr)
+    if not equiv_pattern_reprs:
+        seen_patterns[pattern_repr].append(str(graph) if graph else None)
+        return False
+
+    if graph is None:
+        if skip_duplicates:
+            return True
+        torch._check(
+            False,
+            lambda: f"Duplicate pattern: {pattern_repr} with no graph",
+        )
+
+    new_graph_str = str(graph)
+    for graph_str in equiv_pattern_reprs:
+        if new_graph_str != graph_str:
+            continue
+        if skip_duplicates:
+            return True
+        torch._check(
+            False,
+            lambda: f"Duplicate pattern: {pattern_repr} with duplicated match graph {graph_str} ",
+        )
+    equiv_pattern_reprs.append(new_graph_str)
+    return False
+
+
+def register_replacement(
+    search_fn: SearchFn,
+    replace_fn: ReplaceFn,
+    example_inputs: list[Any] | tuple[Any, ...],
+    trace_fn: TraceFn,
+    pass_dicts: _PassDictsType | Sequence[_PassDictsType],
+    extra_check: Callable[[Match], bool] = _return_true,
+    scalar_workaround: dict[str, float | int] | None = None,
+    exclusive_arg_names: Sequence[str] = (),
+    search_fn_pattern: PatternExpr | None = None,
+    skip_duplicates: bool = False,
+    pattern_name: str | None = None,
+    get_decomp_fn: Callable[..., dict[Any, Callable[..., Any]]] = select_decomp_table,
+) -> bool:
+    """
+    Create a replacement rule based on example functions that get traced
+    to create patterns.  This supports both training and inference when
+    run on a joint forward+backward graph.
+
+    Args:
+        search_fn: traced to give original pattern
+        replace_fn: traced to give replacement graph
+        example_inputs: example inputs for initial trace
+        trace_fn: fwd_only or joint_fwd_bwd
+        pass_dict: dict of passes to register to
+        extra_check: additional check to run on match(using real shapes)
+    """
+    argnames_static = [*inspect.signature(search_fn).parameters.keys()]
+
+    if inspect.ismethod(search_fn):
+        search_fn = _wrap_bound_method(search_fn, argnames_static)
+
+    if inspect.ismethod(replace_fn):
+        replace_argnames = [*inspect.signature(replace_fn).parameters.keys()]
+        replace_fn = _wrap_bound_method(replace_fn, replace_argnames)
+
+    if not isinstance(example_inputs, (list, tuple)):
+        raise TypeError(
+            f"example_inputs must be a list or tuple, got {type(example_inputs)}"
+        )
+
+    def check_fn(match: Match) -> bool:
+        """
+        Often shapes get burned into the pattern, so our initial match ran with
+        `ignore_types=(int, ...)`.
+
+        Recheck the match with the correct shapes.
+        """
+        argnames = list(argnames_static)
+        for name in argnames:
+            if name not in match.kwargs:
+                raise RuntimeError(
+                    f"Not all inputs to pattern found in match.kwargs. Perhaps one "
+                    f"of the inputs is unused? argnames={argnames}, match.kwargs={match.kwargs}"
+                )
+
+        args = list(
+            torch.fx.map_arg(
+                [match.kwargs[name] for name in argnames], lambda n: n.meta["val"]
+            )
+        )
+
+        sym_args: list[torch.SymInt] = []
+        fake_mode = torch._dynamo.utils.detect_fake_mode(args)
+        assert fake_mode is not None
+        with fake_mode:
+            for i, grad in enumerate(requires_grad):
+                if isinstance(args[i], torch.Tensor):
+                    # pyrefly: ignore [missing-attribute]
+                    if grad and is_integer_dtype(args[i].dtype):
+                        return False
+
+                    args[i] = torch.empty_strided(
+                        # pyrefly: ignore [missing-attribute]
+                        args[i].size(),
+                        # pyrefly: ignore [missing-attribute]
+                        args[i].stride(),
+                        # pyrefly: ignore [missing-attribute]
+                        dtype=args[i].dtype,
+                        # pyrefly: ignore [missing-attribute]
+                        device=args[i].device,
+                        requires_grad=grad,
+                    )
+                    # pyrefly: ignore [missing-attribute]
+                    for v in itertools.chain(args[i].shape, args[i].stride()):
+                        if isinstance(v, torch.SymInt) and all(
+                            statically_known_true(v != a) for a in sym_args
+                        ):
+                            sym_args.append(v)
+
+            # If we were given a pre-traced pattern then use that instead of
+            # retracing. Note that this means the pattern has to be independent
+            # of its args.
+            specific_pattern = search_fn_pattern
+
+            if not specific_pattern:
+                if sym_args:
+                    # AOT Autograd and make fx will dedupe symbolic shape size
+                    # accesses of sym ints that appear as inputs
+                    # We don't want the sym_size uses to interfere with pattern matching
+                    # so we provide them as inputs.
+                    # Later, when we actually do the replacement, the symbolic shape
+                    # sizes will get re-traced and added to the graph.
+
+                    def search_fn_new(*args_new: Any, **_: Any) -> Any:
+                        return search_fn(*args_new[len(args_new) - len(args) :])
+
+                    try:
+                        specific_graph = trace_fn(
+                            search_fn_new,
+                            sym_args + args,
+                            get_decomp_fn=get_decomp_fn,
+                        )
+                    except RuntimeError as e:
+                        log_trace_failure(search_fn, e)
+                        return False
+
+                    # correct argnames in the graph
+                    sym_arg_names = []
+                    for i, placeholder in zip(
+                        range(len(sym_args) + len(args)),
+                        specific_graph.graph.nodes,
+                    ):
+                        if i < len(sym_args):
+                            sym_arg_names.append(placeholder.target)
+                            continue
+
+                        with specific_graph.graph.inserting_after(placeholder):
+                            new_node = specific_graph.graph.placeholder(
+                                argnames[i - len(sym_args)]
+                            )
+                            new_node.target = new_node.name
+                            placeholder.replace_all_uses_with(new_node)
+                            specific_graph.graph.erase_node(placeholder)
+
+                    argnames = sym_arg_names + argnames
+                else:
+                    try:
+                        specific_graph = trace_fn(
+                            search_fn, args, get_decomp_fn=get_decomp_fn
+                        )
+                    except RuntimeError as e:
+                        log_trace_failure(search_fn, e)
+                        return False
+
+                specific_pattern = fx_to_pattern(
+                    specific_graph,
+                    argnames=argnames,
+                    exclusive_arg_names=exclusive_arg_names,
+                    scalar_workaround=scalar_workaround,
+                )
+
+            node = match.output_nodes()[0]
+            assert node is not None
+            specific_pattern_match = specific_pattern.match(node)
+
+            if _should_debug_node(node.name):
+                log.warning(
+                    "Specific pattern match: %s%s %s %s",
+                    node,
+                    node.args,
+                    specific_pattern_match,
+                    specific_pattern,
+                )
+
+            if is_match(specific_pattern_match) and extra_check(specific_pattern_match):
+                # trace the pattern using the shapes from the user program
+                match.replacement_graph = trace_fn(
+                    replace_fn, args, get_decomp_fn=get_decomp_fn
+                )
+                if len(match.nodes) == 1:
+                    for n in match.replacement_graph.graph.nodes:
+                        _transfer_meta(
+                            new_meta=n.meta,
+                            old_node=match.nodes[0],
+                            pass_name="replacement",
+                        )
+                return True
+            return False
+
+    def normalize_args(**kwargs: Any) -> list[Any]:
+        args = [kwargs.pop(name) for name in argnames_static]
+        for i in range(1, len(kwargs) + 1):
+            if f"tangents_{i}" not in kwargs:
+                break
+            args.append(kwargs.pop(f"tangents_{i}"))
+        assert not kwargs, f"leftover kwargs: {kwargs!r}"
+        return args
+
+    if trace_fn is joint_fwd_bwd:
+        # If inference mode is enabled during compilation, assume that we don't
+        # want to match on any training graph patterns
+        if torch.is_inference_mode_enabled():
+            return False
+
+    # TODO: Revisit the functionalize_rng_ops for lowmem dropout
+    with functorch_config.patch(functionalize_rng_ops=False):
+        requires_grad: list[bool] = [
+            isinstance(x, torch.Tensor) and x.requires_grad for x in example_inputs
+        ]
+        if search_fn_pattern is None:
+            pattern, gm = gen_pattern_and_search_gm(
+                search_fn,
+                example_inputs,
+                trace_fn,
+                scalar_workaround,
+                exclusive_arg_names,
+                get_decomp_fn=get_decomp_fn,
+            )
+        else:
+            pattern = search_fn_pattern
+            gm = None
+
+        for pattern_matcher_pass in (
+            pass_dicts if isinstance(pass_dicts, Sequence) else [pass_dicts]
+        ):
+            if isinstance(pattern_matcher_pass, PatternMatcherPass):
+                if check_and_add_duplicate_pattern(
+                    pattern,
+                    gm.graph if gm else None,
+                    pattern_matcher_pass.seen_patterns,
+                    skip_duplicates=skip_duplicates,
+                ):
+                    return False
+
+        pattern = ReplacementPatternEntry(
+            pattern=pattern,
+            extra_check=check_fn,
+            normalize_args=normalize_args,
+            pattern_name=pattern_name,
+        )
+        pattern.register(pass_dicts)
+        return pattern.pattern  # type: ignore[return-value]
+
+
+_serialized_patterns: OrderedSet[str] = OrderedSet()
+
+
+def _serialize_pattern(
+    unique_name: str,
+    search_fn: SearchFn,
+    example_inputs: Sequence[Any],
+    trace_fn: TraceFn,
+    scalar_workaround: dict[str, float | int] | None,
+) -> PatternExpr:
+    def get_file_template() -> str:
+        auto_generated_msg = textwrap.dedent(
+            """\
+            # This is an auto-generated file. Please do not modify it by hand.
+            # To re-generate, run:
+            # cd ~/pytorch && python torchgen/fuse/gen_patterns.py
+            """
+        )
+
+        file_template = textwrap.dedent(
+            """\
+            # mypy: ignore-errors
+
+            # noqa: F401, E501
+            {msg}
+            import torch
+            import torch._inductor
+            import operator
+
+            aten = torch.ops.aten
+            prims = torch.ops.prims
+
+            """
+        ).format(msg=auto_generated_msg)
+
+        pattern_matcher_imports = []
+        for name in dir(torch._inductor.pattern_matcher):
+            attr = getattr(torch._inductor.pattern_matcher, name)
+            try:
+                if isinstance(attr, type) and issubclass(
+                    attr, (PatternExpr, _TargetExpr)
+                ):
+                    # pyrefly: ignore [bad-argument-type]
+                    pattern_matcher_imports.append(name)
+            except TypeError:
+                pass
+
+        formatted_imports = ",\n   ".join(pattern_matcher_imports)
+        formatted_imports = f"from torch._inductor.pattern_matcher import (\n   {formatted_imports},\n)\n"
+        return f"{file_template}{formatted_imports}"
+
+    if not SERIALIZED_PATTERN_PATH.is_dir():
+        raise RuntimeError(
+            f"Could not find serialized patterns directory at {SERIALIZED_PATTERN_PATH}"
+        )
+
+    pattern_name = search_fn.__name__
+
+    from torch._functorch import config as functorch_config
+
+    with functorch_config.patch(functionalize_rng_ops=False):
+        pattern = gen_pattern(search_fn, example_inputs, trace_fn, scalar_workaround)
+
+    serialized_pattern = PatternPrettyPrinter.run(pattern, output_name=unique_name)
+    if pattern_name not in _serialized_patterns:
+        write_mode = "w"
+        _serialized_patterns.add(pattern_name)
+    else:
+        write_mode = "a"
+
+    file_template = get_file_template()
+
+    with open(SERIALIZED_PATTERN_PATH / f"{pattern_name}.py", write_mode) as f:
+        if write_mode == "w":
+            f.write(file_template)
+        else:
+            f.write("\n\n")
+        f.write(serialized_pattern)
+        f.write("\n")
+
+    return pattern
+
+
+SERIALIZED_PATTERN_PATH = Path(__file__).parent / "fx_passes" / "serialized_patterns"
+
+# This is the set of serialized patterns that we've registered.  Used by
+# test_serialized_patterns_up_to_date() to ensure the patterns are up
+# to date.
+_known_precompiled_patterns: list[
+    tuple[
+        Any,
+        Iterable[Any],
+        Callable[[Callable[..., Any], Iterable[Any]], torch.fx.GraphModule],
+        Any,
+        PatternExpr,
+    ]
+] = []
+
+
+def gen_register_replacement(
+    unique_name: str,
+    search_fn: SearchFn,
+    replace_fn: ReplaceFn,
+    example_inputs: Iterable[Any],
+    trace_fn: TraceFn,
+    pass_dicts: _PassDictsType | Sequence[_PassDictsType],
+    extra_check: Callable[[Match], bool] = _return_true,
+    scalar_workaround: dict[str, float | int] | None = None,
+    exclusive_arg_names: Sequence[str] = (),
+    skip_duplicates: bool = False,
+    get_decomp_fn: Callable[..., dict[Any, Callable[..., Any]]] = select_decomp_table,
+) -> None:
+    # Make sure the example_inputs is materialized.
+    example_inputs = tuple(example_inputs)
+
+    if "PYTORCH_GEN_PATTERNS" in os.environ:
+        pat = _serialize_pattern(
+            unique_name, search_fn, example_inputs, trace_fn, scalar_workaround
+        )
+    else:
+        pattern_name = search_fn.__name__
+        m = importlib.import_module(
+            f"torch._inductor.fx_passes.serialized_patterns.{pattern_name}"
+        )
+        if not m or not hasattr(m, unique_name):
+            log.warning(
+                "Precompiled pattern %r not found. Run torchgen/fuse/gen_patterns.py.",
+                unique_name,
+            )
+        pat = getattr(m, unique_name)
+
+    for arg in pytree.tree_iter(example_inputs):
+        if isinstance(arg, FakeTensor) and arg.constant is not None:
+            # This can be a problem - small fake tensors (e.g. `tensor(2)`) will
+            # hold onto their original constant value - and by stashing it here
+            # will cause a memory leak if the constant value is on GPU.
+            # Since this is just an optimization we can clear it out.
+            arg.constant = None
+
+    _known_precompiled_patterns.append(
+        (search_fn, example_inputs, trace_fn, scalar_workaround, pat)
+    )
+    register_replacement(
+        search_fn,
+        replace_fn,
+        example_inputs,
+        trace_fn,
+        pass_dicts,
+        extra_check,
+        scalar_workaround,
+        exclusive_arg_names,
+        search_fn_pattern=pat,
+        skip_duplicates=skip_duplicates,
+        pattern_name=unique_name,
+        get_decomp_fn=get_decomp_fn,
+    )
+
+
+@functorch_config.patch(functionalize_rng_ops=False)  # type: ignore[misc]
+def gen_pattern_and_search_gm(
+    search_fn: SearchFn,
+    example_inputs: Sequence[Any],
+    trace_fn: TraceFn,
+    scalar_workaround: dict[str, float | int] | None = None,
+    exclusive_arg_names: Sequence[str] = (),
+    get_decomp_fn: Callable[..., dict[Any, Callable[..., Any]]] = select_decomp_table,
+) -> tuple[PatternExpr, torch.fx.GraphModule]:
+    argnames = [*inspect.signature(search_fn).parameters.keys()]
+
+    if scalar_workaround is None:
+        scalar_workaround = {}
+    flat_inputs = []
+    input_idx = 0  # Positional arguments index
+
+    for argname in argnames:
+        if argname in scalar_workaround:
+            flat_inputs.append(scalar_workaround[argname])
+        else:
+            flat_inputs.append(example_inputs[input_idx])
+            input_idx += 1
+
+    search_gm = trace_fn(search_fn, flat_inputs, get_decomp_fn=get_decomp_fn)
+    return (
+        fx_to_pattern(
+            search_gm,
+            ignore_types=(int, float, list, torch.device, torch.dtype),
+            argnames=argnames,
+            scalar_workaround=scalar_workaround,
+            exclusive_arg_names=exclusive_arg_names,
+        ),
+        search_gm,
+    )
+
+
+def gen_pattern(
+    search_fn: SearchFn,
+    example_inputs: Sequence[Any],
+    trace_fn: TraceFn,
+    scalar_workaround: dict[str, float | int] | None = None,
+    exclusive_arg_names: Sequence[str] = (),
+) -> PatternExpr:
+    return gen_pattern_and_search_gm(
+        search_fn, example_inputs, trace_fn, scalar_workaround, exclusive_arg_names
+    )[0]
+
+
+def register_lowering_pattern(
+    pattern: PatternExpr,
+    extra_check: Callable[[Match], bool] = _return_true,
+    *,
+    pass_dict: _PassDictsType,
+    prepend: bool = False,
+) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
+    """
+    Register an aten to inductor IR replacement pattern.  The decorated
+    function is saved and then called a lowering time allowing direct
+    pattern to inductor IR conversion.
+    """
+
+    def decorator(handler: Callable[..., Any]) -> Callable[..., Any]:
+        assert callable(handler)
+        LoweringPatternEntry(
+            pattern=pattern, extra_check=extra_check, handler=handler
+        ).register(pass_dict, prepend=prepend)
+        handler._inductor_lowering_function = True  # type: ignore[attr-defined]
+        return handler
+
+    return decorator
+
+
+def register_graph_pattern(
+    pattern: PatternExpr,
+    extra_check: Callable[[Match], bool] = _return_true,
+    *,
+    pass_dict: _PassDictsType,
+    prepend: bool = False,
+) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
+    """
+    Register a pattern that runs a function on the FX graph, allowing
+    custom transformation code.
+    """
+
+    def decorator(handler: Callable[..., Any]) -> Callable[..., Any]:
+        assert callable(handler)
+        GraphPatternEntry(
+            pattern=pattern, extra_check=extra_check, handler=handler
+        ).register(pass_dict, prepend=prepend)
+        return handler
+
+    return decorator
+
+
+def is_start_of_fx_graph(graph: torch.fx.Graph, node: torch.fx.Node) -> bool:
+    # first node in the graph
+    return node is next(iter(graph.nodes))
+
+
+# match: copy_, relu_, _set_grad_enabled, manual_seed, _enter_autocast, etc
+# doesn't match: __rshift__, etc
+_mutation_op_re = re.compile(r"(?<!_)(_$|_[.]|(\b|_)(set|enter|exit|seed)(\b|_))(?!_)")
+
+
+def fixme_incorrect_inductor_schema_op(op: torch._ops.OpOverload) -> bool:
+    if op.namespace != "inductor":
+        return False
+
+    # TODO - fix schema
+    # Dont add any more !
+    return op in (
+        torch.ops.inductor.accumulate_grad_.default,
+        torch.ops.inductor.resize_storage_bytes_.default,
+    )
+
+
+def is_mutation_op(node: torch.fx.Node) -> bool:
+    if isinstance(
+        node.target, torch._ops.OpOverload
+    ) and not fixme_incorrect_inductor_schema_op(node.target):
+        return node.target._schema.is_mutable
+    elif isinstance(
+        node.target, torch._higher_order_ops.auto_functionalize.AutoFunctionalized
+    ):
+        return False
+    if node.op == "call_function":
+        assert callable(node.target)
+        if _mutation_op_re.search(node.target.__name__):
+            return True
+    elif node.op == "call_method":
+        assert isinstance(node.target, str)
+        if _mutation_op_re.search(node.target):
+            return True
+    return node.kwargs.get("out") is not None
+
+
+def same_mutation_regions(a: torch.fx.Node, b: torch.fx.Node) -> bool:
+    assert "mutation_region_id" in a.meta
+    assert "mutation_region_id" in b.meta
+    return a.meta["mutation_region_id"] == b.meta["mutation_region_id"]
+
+
+def get_mutation_region_id(graph: torch.fx.Graph, node: torch.fx.Node) -> int:
+    n = node
+    while "mutation_region_id" not in n.meta and not is_start_of_fx_graph(graph, n):
+        n = n.prev
+    mutation_region_id = n.meta.get("mutation_region_id", 0)
+    while n is not node:
+        n = n.next
+        if is_mutation_op(n):
+            mutation_region_id += 1
+        n.meta["mutation_region_id"] = mutation_region_id
+    return mutation_region_id
+
+
+def should_compute_mutation_region_ids(graph: torch.fx.Graph) -> bool:
+    return "mutation_region_id" not in next(iter(graph.nodes)).meta
+
+
+def compute_mutation_region_ids(graph: torch.fx.Graph) -> None:
+    mutation_region_id = 0
+    for nd in graph.nodes:
+        if is_mutation_op(nd):
+            mutation_region_id += 1
+        nd.meta["mutation_region_id"] = mutation_region_id
+
+
+def _wrap_bound_method(fn: Any, argnames: list[str]) -> Any:
+    """
+    Wrap a bound method to remove 'self' from its signature for FX tracing.
+    """
+
+    def wrapper(*args: Any, **kwargs: Any) -> Any:
+        return fn(*args, **kwargs)
+
+    params = [
+        inspect.Parameter(name, inspect.Parameter.POSITIONAL_OR_KEYWORD)
+        for name in argnames
+    ]
+    wrapper.__signature__ = inspect.Signature(params)  # type: ignore[attr-defined]
+    return wrapper
+
+
+class PatternMatcherPass:
+    """
+    Registry of patterns to match and replace in FX graphs.
+    """
+
+    def __init__(
+        self,
+        pass_name: str | None = None,
+        subsystem: str | None = None,
+    ) -> None:
+        super().__init__()
+        self.patterns: defaultdict[
+            tuple[str, torch.fx.node.Target], list[PatternEntry]
+        ] = defaultdict(list)
+        self.pass_name = pass_name
+        self.subsystem = subsystem
+
+        # For a particular generated pattern repr, store all of the str representations
+        # of the graph used to generate them. Because we ignore certain patterns
+        # in searching, but not in matching, use the graph to distinguish if two equivalent
+        # searches are actually different.
+        self.seen_patterns: dict[str, list[str | None]] = defaultdict(list)
+
+    def __getitem__(self, item: tuple[str, torch.fx.node.Target]) -> list[PatternEntry]:
+        return self.patterns[item]
+
+    def apply(self, gm: torch.fx.GraphModule | torch.fx.Graph) -> int:
+        """Apply all registered patterns to the graph, returning the number of matches."""
+        if not self.patterns:
+            return 0
+        if isinstance(gm, torch.fx.GraphModule):
+            graph = gm.graph
+        elif isinstance(gm, torch.fx.Graph):
+            graph = gm
+            gm = graph.owning_module
+        else:
+            raise RuntimeError(
+                f"The input to PatternMatcherPass must be a GraphModule or a Graph, but got {type(gm)}"
+            )
+        if should_compute_mutation_region_ids(graph):
+            compute_mutation_region_ids(graph)
+        get_mutation_region_id_partial = functools.partial(
+            get_mutation_region_id, graph
+        )
+        count = 0
+        nodes = []
+        has_call_module = False
+        for op, target in self.patterns:
+            if op == "call_module":
+                has_call_module = True
+            else:
+                nodes.append(graph.find_nodes(op=op, target=target, sort=False))
+        if has_call_module:
+            nodes.append(graph.find_nodes(op="call_module", sort=False))
+        pass_name = self.pass_name if self.pass_name is not None else "pattern_matcher"
+        assert isinstance(gm, torch.fx.GraphModule)
+        with GraphTransformObserver(gm, pass_name, self.subsystem):
+            for node in sorted(itertools.chain.from_iterable(nodes), reverse=True):
+                target = extract_target(node)
+                if node.op == "call_module":
+                    if (node.op, target) not in self.patterns:
+                        continue
+
+                # conservatively not applying pattern for cpu input,
+                # since some of the patterns induce codegen and split nodes.
+                # Note: we will only skip cpu compute if disable_cpp_codegen=True
+                if fallback_node_due_to_unsupported_type(node, allow_cpu_inputs=False):
+                    continue
+
+                for entry in self.patterns[(node.op, target)]:
+                    if node._erased:
+                        break
+                    m = entry.pattern.match(node)
+                    # pattern match crosses mutation barrier - discard
+                    if (
+                        is_match(m)
+                        and len(
+                            OrderedSet(map(get_mutation_region_id_partial, m.nodes))
+                        )
+                        != 1
+                    ):
+                        continue
+                    # pattern match crosses stream boundary - discard
+                    if (
+                        is_match(m)
+                        and len(
+                            OrderedSet(
+                                n.meta.get("custom", {}).get("stream", 0)
+                                for n in m.nodes
+                            )
+                        )
+                        != 1
+                    ):
+                        continue
+                    if _should_debug_node(node.name):
+                        log.warning("%s%s %s %s", node, node.args, m, entry.pattern)
+
+                    if is_match(m) and guard_or_false(entry.extra_check(m)):
+                        count += 1
+                        entry.apply(m, graph, node)
+                        counters[backend]["pattern_matcher_count"] += 1
+                        counters[backend]["pattern_matcher_nodes"] += len(m.nodes)
+
+                        # Track per-pattern counts when debug mode is active
+                        if os.environ.get("TORCHINDUCTOR_PATTERN_MATCH_DEBUG"):
+                            if getattr(entry, "pattern_name", None):
+                                pattern_name = entry.pattern_name
+                            else:
+                                # Fallback: use pattern class name + operation target
+                                pattern_class = entry.pattern.__class__.__name__
+                                target = str(node.target)
+                                pattern_name = f"{pattern_class}_{target}"
+
+                            pattern_key = f"{backend}_pattern_matcher_per_pattern"
+                            counters[pattern_key][pattern_name] += 1
+        return count
+
+    def clear(self) -> None:
+        self.patterns.clear()
+
+
+def _not_implemented(*args: Any, **kwargs: Any) -> NoReturn:
+    raise NotImplementedError
+
+
+def fx_to_pattern(
+    gm: torch.fx.GraphModule | torch.fx.Graph,
+    ignore_types: Sequence[type[Any]] = (),
+    argnames: Sequence[str] = (),
+    scalar_workaround: dict[str, float | int] | None = None,
+    exclusive_arg_names: Sequence[str] = (),
+) -> PatternExpr:
+    """
+    Convert an FX graph into a PatternExpr.  This is useful for simple
+    patterns that can only match single functions and fixed-length lists.
+    """
+    # scalar_workaround is a hack to capture dropout_p
+    # see https://github.com/pytorch/pytorch/issues/97894
+    scalar_workaround = scalar_workaround or {}
+    inv_scalar_workaround = {v: k for k, v in scalar_workaround.items()}
+    assert len(inv_scalar_workaround) == len(scalar_workaround)
+
+    def process_arg(
+        x: T, ignore_types_override: Sequence[type[Any]] | None = None
+    ) -> T | KeywordArg | Ignored:
+        current_ignore_types = (
+            ignore_types_override if ignore_types_override is not None else ignore_types
+        )
+        if isinstance(x, (float, int)) and x in inv_scalar_workaround:
+            return KeywordArg(inv_scalar_workaround[x])
+        if type(x) in current_ignore_types:
+            return Ignored()
+        if isinstance(x, list) and all(isinstance(y, Ignored) for y in x) and x:
+            return Ignored()
+        return x
+
+    argnum = itertools.count()
+
+    class Converter(torch.fx.Interpreter):
+        # pyrefly: ignore [bad-override]
+        call_method = _not_implemented
+        # pyrefly: ignore [bad-override]
+        call_module = _not_implemented
+        # pyrefly: ignore [bad-override]
+        get_attr = _not_implemented
+
+        # pyrefly: ignore [bad-override]
+        def placeholder(
+            self,
+            target: str,  # type: ignore[override]
+            args: Sequence[Any],
+            kwargs: Mapping[str, Any],
+        ) -> ExclusiveKeywordArg | KeywordArg:
+            n = next(argnum)
+            if n < len(argnames):
+                name = argnames[n]
+            elif argnames:
+                assert target.startswith("tangent")
+                name = target
+            else:
+                target = re.sub(r"_\d+$", "", target)  # de-mangle arg name
+                name = target
+            if name in exclusive_arg_names:
+                return ExclusiveKeywordArg(name)
+            else:
+                return KeywordArg(name)
+
+        # pyrefly: ignore [bad-override]
+        def call_function(
+            self,
+            target: str,  # type: ignore[override]
+            args: Sequence[Any],
+            kwargs: Mapping[str, Any],
+        ) -> PatternExpr:
+            process_arg_fn = process_arg
+            # Indexing is critical for matching getitem nodes, so we can't ignore int args here
+            if target is operator.getitem:
+
+                def process_arg_fn_impl(
+                    x: T,
+                    ignore_types_override: Sequence[type[Any]] | None = tuple(
+                        t for t in ignore_types if t is not int
+                    ),
+                ) -> T | KeywordArg | Ignored:
+                    return process_arg(x, ignore_types_override)
+
+                process_arg_fn = process_arg_fn_impl
+
+            args, kwargs = pytree.tree_map(process_arg_fn, (args, kwargs))
+            if list in ignore_types:
+                # Handle a burned in tensor size which are now [Ignored(), Ignored(), ...]
+                args = [process_arg_fn(a) for a in args]
+                kwargs = {k: process_arg_fn(a) for k, a in kwargs.items()}
+            return CallFunction(target, *args, **kwargs)
+
+        def run_node(self, n: torch.fx.Node) -> Any:
+            rv = super().run_node(n)
+            if n.op == "output" and isinstance(rv, tuple):
+                args = n.args[0]
+                assert isinstance(args, Collection)
+                assert len(rv) == len(args)
+                for r, arg in zip(rv, args):
+                    # pyrefly: ignore [missing-attribute]
+                    r.users = len(arg.users)
+            else:
+                rv.users = len(n.users)
+            return rv
+
+    assert isinstance(gm, torch.fx.GraphModule)
+    pattern = Converter(gm).run()
+    if not isinstance(pattern, PatternExpr):
+        return MultiOutputPattern(pytree.tree_leaves(pattern))
+    return pattern
+
+
+@torch.no_grad()
+def fwd_only(
+    fn: Callable[..., Any],
+    args: Sequence[Any],
+    *,
+    run_functional_passes: bool = True,
+    get_decomp_fn: Callable[..., Any] = select_decomp_table,
+) -> torch.fx.GraphModule:
+    """Build a normalized inference graph, for use with fx_to_pattern"""
+    # TODO - look into using aot autograd, asserting no mutating ops here
+    with enable_python_dispatcher(), preserve_node_meta():
+        gm = make_fx(fn, get_decomp_fn(), tracing_mode="real")(*args)
+
+    from .fx_passes.post_grad import remove_noop_ops
+
+    if run_functional_passes:
+        remove_noop_ops(gm.graph)
+
+        # NOTE: applying early_patterns to user patterns cause
+        # duplicate patterns being found in vllm. Check
+        # https://github.com/pytorch/pytorch/pull/170649#issuecomment-3693427775
+        # for more details.
+        # from .fx_passes.joint_graph import early_patterns
+        # early_patterns.apply(gm.graph)
+
+        gm.graph.eliminate_dead_code()
+
+    gm.recompile()
+    return gm
+
+
+@torch.enable_grad()
+def joint_fwd_bwd(
+    fn: Callable[..., Any],
+    args: Sequence[Any],
+    *,
+    get_decomp_fn: Callable[..., Any] = select_decomp_table,
+) -> torch.fx.GraphModule:
+    """Build a normalized training graph, for use with fx_to_pattern"""
+    gm: torch.fx.GraphModule | None = None
+
+    def record_joint_graph(
+        joint_graph: torch.fx.GraphModule, inputs: Sequence[Any], **kwargs: Any
+    ) -> tuple[torch.fx.GraphModule, torch.fx.GraphModule]:
+        nonlocal gm
+        assert not gm
+        gm = clone_graph(joint_graph)
+        return default_partition(joint_graph, inputs, **kwargs)
+
+    with torch._guards.tracing(None):
+        aot_function(
+            fn,
+            # pyrefly: ignore[bad-argument-type]
+            lambda gm, example_inputs: make_boxed_func(gm),
+            partition_fn=record_joint_graph,
+            decompositions=get_decomp_fn(),
+            keep_inference_input_mutations=True,
+            enable_log=False,
+        )(*args)
+    assert gm
+
+    from .fx_passes.post_grad import remove_noop_ops
+
+    remove_noop_ops(gm.graph)
+
+    from .fx_passes.joint_graph import early_patterns
+
+    early_patterns.apply(gm.graph)
+
+    # remove in/out specs
+    gm.graph._codegen = torch.fx.graph.CodeGen()
+    gm.graph.eliminate_dead_code()
+    gm.recompile()
+    return gm
+
+
+def _args(n: torch.fx.Node) -> list[torch.fx.node.Argument]:
+    args: list[torch.fx.node.Argument] = []
+    torch.fx.map_arg((n.args, n.kwargs), args.append)
+    return args
+
+
+def stable_topological_sort(graph: torch.fx.Graph) -> None:
+    # Nodes are in exactly one of these three collections:
+
+    # - Nodes in `pending` are waiting to be processed (in reverse order):
+    pending = list(reversed(graph.nodes))
+
+    # - Nodes in `ready` have been processed and are already in the correct
+    #   order.
+    ready = OrderedSet[torch.fx.Node]()
+
+    # - `waiting` is a mapping from a dependency to nodes which depend on that
+    #   dependency.
+    waiting = defaultdict(list)
+
+    # The cursor indicates the last processed node so we can add new nodes
+    # after it.
+    cursor = None
+    while pending:
+        node = pending.pop()
+        waiting_for = [x for x in _args(node) if x not in ready]
+        if waiting_for:
+            # We have unprocessed input nodes. Might as well wait for the last
+            # arg so an already sorted list will only recheck this node once.
+            waiting[waiting_for[-1]].append(node)
+        else:
+            ready.add(node)
+            if cursor and cursor.next is not node:
+                cursor.append(node)
+            cursor = node
+            # Mark the nodes that have been waiting for this node to finish as
+            # ready to check again.
+            pending.extend(reversed(waiting.pop(node, ())))
+
+    assert not waiting and len(ready) == len(graph.nodes)
+
+
+def init_once_fakemode(fn: Callable[..., Any]) -> Callable[..., Any]:
+    """Wrapper around lazy init functions in fx_passes/"""
+
+    _fn_params = inspect.signature(fn).parameters
+
+    @functools.cache
+    @functools.wraps(fn)
+    def lazy_init(
+        input_device: Any | None = None,
+        get_decomp_fn: Callable[
+            ..., dict[Any, Callable[..., Any]]
+        ] = select_decomp_table,
+    ) -> Any:
+        counters_ref = counters[backend].copy()
+
+        kwargs: dict[str, Any] = {}
+        if "input_device" in _fn_params:
+            kwargs["input_device"] = input_device
+        if "get_decomp_fn" in _fn_params:
+            kwargs["get_decomp_fn"] = get_decomp_fn
+
+        with torch._guards.tracing(None), unset_fake_temporarily(), FakeTensorMode():
+            result = fn(**kwargs)
+
+        # clear view matches encountered during tracing
+        counters[backend] = counters_ref
+
+        return result
+
+    return lazy_init
+
+
+def config_flag(name: str) -> Callable[[Match], Any]:
+    """Function for extra_check to put pass behind a flag"""
+
+    def flag_check(match: Match) -> Any:
+        return getattr(config, name)
+
+    return flag_check
+
+
+def clone_graph(input_graph: torch.fx.GraphModule) -> torch.fx.GraphModule:
+    class CopyGraph(Transformer):
+        def run_node(self, old_node: torch.fx.Node) -> torch.fx.Node:
+            new_node = super().run_node(old_node)
+            if isinstance(new_node, torch.fx.Proxy):
+                new_node.node.meta.update(old_node.meta)
+                new_node.node.name = self.new_graph._graph_namespace.create_name(
+                    old_node.name, None
+                )
+
+            return new_node
+
+    return CopyGraph(input_graph).transform()
+
+
+# TODO: remove in follow up diff, used internally
+_seen_patterns: OrderedSet[str] = OrderedSet()
+
+
+def get_arg_value(
+    node: torch.fx.Node, arg_number: int, kwarg_name: str | None = None
+) -> Any:
+    if len(node.args) > arg_number:
+        return node.args[arg_number]
+    elif kwarg_name is None:
+        return None
+    else:
+        return node.kwargs.get(kwarg_name)
+
+
+def filter_nodes(nodes: Iterable[torch.fx.Node], fn: Any) -> list[torch.fx.Node]:
+    fns = [fn]
+    if isinstance(fn, torch._ops.OpOverloadPacket):
+        fns.extend([getattr(fn, overload) for overload in fn.overloads()])
+
+    return [node for node in nodes if node.target in fns]
+
+
+def extract_target(node: torch.fx.Node) -> torch.fx.node.Target:
+    """For call_function and call_method, we directly use the target function;
+    For call_module, the target is string, and we treat the module class
+     as a function.
+    """
+    if node.op == "call_module":
+        assert isinstance(node.target, str)
+        return _get_attr(node.graph.owning_module, node.target).__class__
+    return node.target

outputs/audit_venv/lib/python3.11/site-packages/torch/_inductor/quantized_lowerings.py

@@ -0,0 +1,171 @@
+import logging
+from typing import Any
+
+import torch
+from torch._inductor.kernel.mm_common import mm_args
+
+from . import config, lowering
+from .codegen.cpp_gemm_template import CppGemmTemplate, CppWoqInt4GemmTemplate
+from .codegen.cpp_utils import create_epilogue_with_attr
+from .lowering import expand, register_lowering
+from .mkldnn_ir import WeightInt4PackMatmul
+from .select_algorithm import (
+    autotune_select_algorithm,
+    ExternKernelChoice,
+    realize_inputs,
+)
+from .utils import use_aten_gemm_kernels, use_cpp_gemm_template
+from .virtualized import V
+
+
+log = logging.getLogger(__name__)
+
+aten__weight_int8pack_mm = ExternKernelChoice(
+    torch._weight_int8pack_mm, "at::_weight_int8pack_mm", has_out_variant=False
+)
+
+aten__weight_int4pack_mm_cpu = ExternKernelChoice(
+    torch.ops.quantized.int4mm_packed_weight_cpu,
+    "at::native::_weight_int4pack_mm_cpu_tensor",
+    has_out_variant=False,
+    kernel_creator=WeightInt4PackMatmul.create,
+)
+
+quantized = torch.ops.quantized
+_quantized = torch.ops._quantized
+aten = torch.ops.aten
+
+
+def register_quantized_ops() -> None:
+    lowering.add_needs_realized_inputs(
+        [
+            quantized.max_pool2d,
+            _quantized.wrapped_fbgemm_pack_gemm_matrix_fp16,
+            _quantized.wrapped_fbgemm_linear_fp16_weight,
+        ]
+    )
+    lowering.make_fallback(quantized.max_pool2d)
+    lowering.make_fallback(_quantized.wrapped_fbgemm_pack_gemm_matrix_fp16)
+    lowering.make_fallback(_quantized.wrapped_fbgemm_linear_fp16_weight)
+
+
+def register_woq_mm_ops() -> None:
+    @register_lowering(aten._weight_int8pack_mm, type_promotion_kind=None)  # type: ignore[misc]
+    def int8pack_mm(
+        input: torch.Tensor,
+        weight: torch.Tensor,
+        scale: torch.Tensor,
+        *,
+        layout: Any = None,
+    ) -> Any:
+        _, _, _, layout, mat1, mat2 = mm_args(
+            input, weight, layout=layout, mat2_transposed=True
+        )
+        assert (
+            mat1.get_dtype() in [torch.bfloat16, torch.float16, torch.float]
+            and mat2.get_dtype() == torch.int8
+        )
+        aten_layout = layout
+
+        # options to tune from
+        choices = (
+            [aten__weight_int8pack_mm.bind((mat1, mat2, scale), aten_layout)]
+            if use_aten_gemm_kernels()
+            else []
+        )
+
+        # scale is applied as an epilogue, and the scale tensor is expanded (with a view op)
+        # for broadcasting, as it's 1D.
+        def _mul_epilogue(buf: torch.Tensor) -> Any:
+            return create_epilogue_with_attr(
+                buf, "mul", other=realize_inputs(expand(scale, layout.size))
+            )
+
+        if use_cpp_gemm_template(aten_layout, mat1, mat2, mat2_transposed=True):
+            CppGemmTemplate.add_choices(
+                choices,
+                aten_layout,
+                [mat1, mat2, scale],
+                trans_w=True,
+                epilogue_creator=_mul_epilogue,  # type: ignore[arg-type]
+            )
+
+        node, _ = autotune_select_algorithm(
+            "_weight_int8pack_mm", choices, [mat1, mat2, scale], aten_layout
+        )
+        return node
+
+    @register_lowering(aten._weight_int4pack_mm_for_cpu, type_promotion_kind=None)  # type: ignore[misc]
+    def int4pack_mm_cpu(
+        input: torch.Tensor,
+        weight: torch.Tensor,
+        qGroupSize: int,
+        qScaleAndZeros: torch.Tensor,
+        *,
+        layout: Any = None,
+    ) -> Any:
+        _, _, _, layout, mat1, mat2 = mm_args(
+            input, weight, layout=layout, use_4x2_dim=True, mat2_transposed=True
+        )
+        assert (
+            mat1.get_dtype() in [torch.bfloat16, torch.float16, torch.float]
+            and mat2.get_dtype() == torch.uint8
+        )
+        group_size = V.graph.add_tensor_constant(
+            torch.tensor(qGroupSize, dtype=torch.int64), name=None
+        )
+        aten_layout = layout
+
+        # options to tune from
+        choices = (
+            [
+                aten__weight_int4pack_mm_cpu.bind(
+                    (mat1, mat2, group_size, qScaleAndZeros), aten_layout
+                )
+            ]
+            if use_aten_gemm_kernels()
+            else []
+        )
+        if (
+            (config.max_autotune or config.max_autotune_gemm)
+            and use_cpp_gemm_template(
+                aten_layout,
+                mat1,
+                mat2,
+                mat2_transposed=True,
+                is_woq_int4=True,
+                q_group_size=qGroupSize,
+            )
+            and mat2.get_layout().is_contiguous()
+        ):
+            # pyrefly: ignore [bad-specialization, missing-attribute, not-a-type]
+            CppWoqInt4GemmTemplate[qGroupSize].add_choices(
+                choices,
+                aten_layout,
+                [mat1, mat2, group_size, qScaleAndZeros],
+            )
+
+        # define functions to generate example inputs for weight and group size
+        # otherwise, autotuner generates example inputs of all zeros for them
+        def get_example_weight(x: torch._inductor.ir.IRNode) -> torch.Tensor:
+            assert x.get_layout().is_contiguous()
+            shape = x.get_size()
+            device = x.get_device()
+            return torch.randint(0, 255, shape, dtype=torch.uint8, device=device)
+
+        input_gen_fns = {
+            1: get_example_weight,  # packed weight
+            2: lambda x: V.graph.constants[x.get_name()],  # group size
+        }
+
+        node, _ = autotune_select_algorithm(
+            "_weight_int4pack_mm_for_cpu",
+            choices,
+            [mat1, mat2, group_size, qScaleAndZeros],
+            aten_layout,
+            input_gen_fns=input_gen_fns,
+        )
+        return node
+
+    lowering.make_fallback(aten._dyn_quant_matmul_4bit)
+    lowering.make_fallback(aten._dyn_quant_pack_4bit_weight)

outputs/audit_venv/lib/python3.11/site-packages/torch/_inductor/remote_cache.py

@@ -0,0 +1,432 @@
+from __future__ import annotations
+
+import atexit
+import collections
+import dataclasses
+import functools
+import json
+import logging
+import os
+import sys
+import typing
+from abc import abstractmethod
+from typing import Any, Generic, TypeAlias, TypeVar, Union
+from typing_extensions import override
+
+from torch._dynamo.utils import dynamo_timed
+from torch._inductor import config
+from torch.monitor import _WaitCounter
+
+
+if typing.TYPE_CHECKING:
+    from collections.abc import Callable
+
+
+try:
+    import redis
+except ImportError:
+    redis = None  # type: ignore[assignment]
+
+
+log = logging.getLogger(__name__)
+
+
+if config.is_fbcode():
+    from rfe.scubadata.scubadata_py3 import (  # type: ignore[import-not-found]
+        Sample as Sample_,
+    )
+
+    Sample: TypeAlias = Sample_
+else:
+    Sample: TypeAlias = type[object]  # type: ignore[misc,no-redef]
+
+
+_T = TypeVar("_T")
+_U = TypeVar("_U")
+
+
+remote_fx_cache_get_timed = functools.partial(
+    dynamo_timed,
+    "FbRemoteFxGraphCache.get",
+    phase_name="remote_fx_graph_cache_get",
+    log_pt2_compile_event=False,
+    dynamo_compile_column_us="remote_fx_graph_cache_get_time_us",
+    log_waitcounter=True,
+)
+remote_fx_cache_put_timed = functools.partial(
+    dynamo_timed,
+    "FbRemoteFxGraphCache.put",
+    phase_name="remote_fx_graph_cache_put",
+    log_pt2_compile_event=False,
+    dynamo_compile_column_us="remote_fx_graph_cache_put_time_us",
+    log_waitcounter=True,
+)
+
+
+class RemoteCacheBackend(Generic[_T]):
+    """
+    A backend implementation for accessing a remote/distributed cache.  Only
+    works with bytes in/out.  For structured data use a RemoteCache.
+    """
+
+    def __init__(self) -> None:
+        self._name = f"backend:{type(self).__name__}"
+
+    @abstractmethod
+    def _get(self, key: str) -> _T | None:
+        pass
+
+    @abstractmethod
+    def _put(self, key: str, data: _T) -> None:
+        pass
+
+    def get(self, key: str) -> _T | None:
+        try:
+            value = self._get(key)
+            cache_stats.get(self._name, value)
+        except Exception:
+            cache_stats.exception(self._name)
+            raise
+        return value
+
+    def put(self, key: str, data: _T) -> None:
+        try:
+            self._put(key, data)
+            cache_stats.put(self._name)
+        except Exception:
+            cache_stats.exception(self._name)
+            raise
+
+
+# Serde that encodes from _T to _U and decodes from _U to _T.
+class RemoteCacheSerde(Generic[_T, _U]):
+    @abstractmethod
+    def encode(self, data: _T) -> _U:
+        pass
+
+    @abstractmethod
+    def decode(self, data: _U) -> _T:
+        pass
+
+
+JsonDataTy: TypeAlias = Union[  # noqa: UP007
+    int, float, str, bool, dict[str, "JsonDataTy"], list["JsonDataTy"], None
+]
+
+
+class RemoteCacheJsonSerde(RemoteCacheSerde[JsonDataTy, bytes]):
+    def encode(self, data: JsonDataTy) -> bytes:
+        return bytes(json.dumps(data), "ascii")
+
+    def decode(self, data: bytes) -> JsonDataTy:
+        return json.loads(data)
+
+
+class RemoteCachePassthroughSerde(RemoteCacheSerde[_T, _T]):
+    def encode(self, data: _T) -> _T:
+        return data
+
+    def decode(self, data: _T) -> _T:
+        return data
+
+
+# This class is the top of a RemoteCache. A RemoteCache is fundamentally made of
+# three parts:
+#
+# 1. The controller (this class).
+# 2. A serializer/deserializer (instance of RemoteCacheSerde).
+# 3. A backend (instance of RemoteCacheBackend).
+#
+# To write (`put`), the RemoteCache takes data, uses the RemoteCacheSerde to
+# convert it for the backend and passes it to the backend.
+#
+# Conversely when reading (`get`), the RemoteCache takes data from the backend,
+# uses the RemoteCacheSerde to convert it and returns it.
+#
+# The RemoteCacheBackend is generic on _U - which is the type of data the
+# backend can directly cache (usually `bytes`).
+#
+# The RemoteCacheSerde is responsible for converting between _T (the type of
+# data the RemoteCache accepts in `put` and returns in `get`) and _U.
+#
+# When instantiating a RemoteCache you should override, not directly create a
+# RemoteCache. The reason is that when logging cache use (`TORCH_LOGS=cache`) we
+# use the concrete type of the RemoteCache as the reported cache. See
+# RemoteFxGraphCache below as an example.
+class RemoteCache(Generic[_T]):
+    backend_override_cls: Callable[[], RemoteCacheBackend[Any]] | None = None
+
+    def __init__(
+        self, backend: RemoteCacheBackend[_U], serde: RemoteCacheSerde[_T, _U]
+    ) -> None:
+        # Support for testing to mock out the backend on a class-by-class basis.
+        if (override_cls := self.__class__.backend_override_cls) is not None:
+            self.backend = override_cls()
+        else:
+            self.backend = backend
+        # pyrefly: ignore [invalid-type-var]
+        self.serde = serde
+
+    # See if the cache contains `key`. Returns `None` if the value is not
+    # present in the cache.
+    def get(self, key: str) -> _T | None:
+        with _WaitCounter("pytorch.remote_cache.get").guard():
+            sample = self._create_sample()
+            try:
+                result = self._get(key, sample)
+                cache_stats.get(type(self).__name__, result)
+            except Exception as e:
+                cache_stats.exception(type(self).__name__)
+                if sample:
+                    sample.fail_reason = str(e)
+                raise
+            finally:
+                self._log_sample(sample)
+            return result
+
+    # Add `value` to the cache with the key `key`. Note that `None` is not a
+    # valid value even if _T supports it (because you can't tell the difference
+    # between `None` and a missing cache entry).
+    def put(self, key: str, value: _T) -> None:
+        with _WaitCounter("pytorch.remote_cache.put").guard():
+            assert value is not None
+            sample = self._create_sample()
+            try:
+                self._put(key, value, sample)
+                cache_stats.put(type(self).__name__)
+            except Exception as e:
+                cache_stats.exception(type(self).__name__)
+                if sample:
+                    sample.fail_reason = str(e)
+                raise
+            finally:
+                self._log_sample(sample)
+
+    # Used to convert data from the cache into structured data.
+    def _decode(self, data: _U, sample: Sample | None) -> _T:  # type: ignore[override]
+        return self.serde.decode(data)  # type: ignore[arg-type]
+
+    # Used to convert structured data into data for the cache.
+    def _encode(self, value: _T, sample: Sample | None) -> object:  # returns _U
+        return self.serde.encode(value)
+
+    # Get structured data from the cache.
+    # Separate from `get` so that it can be overridden.
+    def _get(self, key: str, sample: Sample | None) -> _T | None:
+        if data := self._backend_get(key):
+            return self._decode(data, sample)
+        return None
+
+    # Get unstructured data from the cache.
+    # Separate from `get` so that it can be overridden.
+    # Returns _U - but we aren't actually generic on _U
+    def _backend_get(self, key: str) -> object:
+        return self.backend.get(key)
+
+    # Put structured data into the cache.
+    # Separate from `put` so that it can be overridden.
+    def _put(self, key: str, value: _T, sample: Sample | None) -> None:
+        data = self._encode(value, sample)
+        self._backend_put(key, data)
+
+    # Put unstructured data into the cache.
+    # Separate from `put` so that it can be overridden.
+    # Takes data: _U - but we aren't actually generic on _U
+    def _backend_put(self, key: str, data: object) -> None:
+        self.backend.put(key, data)
+
+    # Create a logging Sample - used with internal loggers to monitor cache
+    # effectiveness.
+    def _create_sample(self) -> Sample | None:
+        return None
+
+    # Write the logging Sample to the logger.
+    def _log_sample(self, sample: Sample | None) -> None:
+        pass
+
+
+class RedisRemoteCacheBackend(RemoteCacheBackend[bytes]):
+    """
+    A Redis implementation of a remote/distributed cache.
+    """
+
+    # pyrefly: ignore [missing-attribute]
+    _redis: redis.Redis | None = None
+
+    def __init__(self, cache_id: str) -> None:
+        super().__init__()
+        if not redis:
+            raise RuntimeError("redis not available but required for remote cache")
+
+        if "TORCHINDUCTOR_REDIS_URL" in os.environ:
+            self._redis = redis.Redis.from_url(os.environ["TORCHINDUCTOR_REDIS_URL"])
+        else:
+            self._redis = redis.Redis(
+                host=os.environ.get("TORCHINDUCTOR_REDIS_HOST", "localhost"),
+                port=int(os.environ.get("TORCHINDUCTOR_REDIS_PORT", 6379)),
+            )
+
+    @override
+    def _get(self, key: str) -> bytes | None:
+        if not self._redis:
+            # Either redis wasn't found or we already had some trouble...
+            return None
+
+        try:
+            # pyrefly: ignore [missing-attribute]
+            value = self._redis.get(key)
+        # pyrefly: ignore [missing-attribute]
+        except redis.exceptions.ConnectionError:
+            # Redis is lazy and doesn't actually attempt to connect until the
+            # first use. Mark is as unavailable now.
+            self._redis = None
+            return None
+
+        # In theory redis.get() can return an Awaitable as well...
+        assert value is None or isinstance(value, bytes)
+        return value
+
+    @override
+    def _put(self, key: str, data: bytes) -> None:
+        if not self._redis:
+            # Either redis wasn't found or we already had some trouble...
+            return
+
+        try:
+            # pyrefly: ignore [missing-attribute]
+            self._redis.set(key, data)
+        # pyrefly: ignore [missing-attribute]
+        except redis.exceptions.ConnectionError:
+            # Redis is lazy and doesn't actually attempt to connect until the
+            # first use. Mark is as unavailable now.
+            self._redis = None
+
+
+class RedisRemoteCache(RemoteCache[JsonDataTy]):
+    def __init__(self, cache_id: str) -> None:
+        # Special test handling: If we're just going to override the backend
+        # anyway don't require redis
+        if self.__class__.backend_override_cls:
+            # This is totally bogus but it works for now...
+            backend = typing.cast(RemoteCacheBackend[bytes], None)
+        else:
+            backend = RedisRemoteCacheBackend(cache_id)
+        serde = RemoteCacheJsonSerde()
+        super().__init__(backend, serde)
+        version = 1  # consistency between various types of keys
+        self._key_fmt = f"pt2:{cache_id}::{{key}}:c{version}"
+
+    def _get_key(self, key: str) -> str:
+        return self._key_fmt.format(key=key)
+
+    @override
+    def _get(self, key: str, sample: Sample | None) -> JsonDataTy | None:
+        key = self._get_key(key)
+        return super()._get(key, sample)
+
+    @override
+    def _put(self, key: str, value: JsonDataTy, sample: Sample | None) -> None:
+        key = self._get_key(key)
+        super()._put(key, value, sample)
+
+
+class RemoteAutotuneCache(RedisRemoteCache):
+    pass
+
+
+class RemoteBundledAutotuneCache(RedisRemoteCache):
+    pass
+
+
+class RemoteFxGraphCache(RedisRemoteCache):
+    pass
+
+
+class RemoteAOTAutogradCache(RedisRemoteCache):
+    pass
+
+
+class RemoteDynamoPGOCache(RedisRemoteCache):
+    pass
+
+
+def create_cache(
+    key: str,
+    is_fbcode: bool,
+    fb_cache_cls: str,
+    oss_cache_cls: str,
+) -> RemoteCache[JsonDataTy] | None:
+    try:
+        if is_fbcode:
+            import torch._inductor.fb.remote_cache
+
+            cache_cls = getattr(torch._inductor.fb.remote_cache, fb_cache_cls)
+            return cache_cls(key)
+        else:
+            this_module = sys.modules[__name__]
+
+            cache_cls = getattr(this_module, oss_cache_cls)
+            return cache_cls(key)
+
+    except Exception:
+        log.warning("Unable to create a remote cache", exc_info=True)
+        return None
+
+
+# Some simple stat capture
+@dataclasses.dataclass
+class _CacheStat:
+    miss: int = 0
+    hit: int = 0
+    put: int = 0
+    exception: int = 0
+
+    def __str__(self) -> str:
+        return f"{{hit: {self.hit}, miss: {self.miss}, put: {self.put}, exception: {self.exception}}}"
+
+
+class _CacheStats:
+    _stats: dict[str, _CacheStat]
+
+    def __init__(self) -> None:
+        self._stats = collections.defaultdict(_CacheStat)
+
+    def miss(self, name: str, count: int = 1) -> None:
+        self._stats[name].miss += count
+
+    def hit(self, name: str, count: int = 1) -> None:
+        self._stats[name].hit += count
+
+    def get(self, name: str, value: object | None) -> None:
+        if value is None:
+            self.miss(name)
+        else:
+            self.hit(name)
+
+    def put(self, name: str, count: int = 1) -> None:
+        self._stats[name].put += count
+
+    def exception(self, name: str, count: int = 1) -> None:
+        self._stats[name].exception += count
+
+
+cache_stats = _CacheStats()
+
+
+@atexit.register
+def dump_cache_stats() -> None:
+    if not log.isEnabledFor(logging.INFO):
+        return
+
+    import io
+
+    out = io.StringIO()
+
+    if not cache_stats._stats:
+        print(" None", file=out)
+    else:
+        print(file=out)
+        for k, v in sorted(cache_stats._stats.items()):
+            print(f"  {k}: {v}", file=out)
+
+    log.info("Cache Metrics:%s", out.getvalue())

outputs/audit_venv/lib/python3.11/site-packages/torch/_inductor/remote_gemm_autotune_cache.py

@@ -0,0 +1,20 @@
+import asyncio
+from typing import TypeVar
+
+import torch._inductor.config as config
+from torch._inductor import ir
+
+
+_T = TypeVar("_T")
+
+
+def gen_best_config(mat1: ir.StorageBox, mat2: ir.StorageBox) -> asyncio.Task[_T]:
+    """
+    Generate the best GEMM autotune config for the given matrices.
+    """
+    if config.is_fbcode():
+        from torch._inductor.fb.remote_gemm_autotune_cache import gen_best_config
+
+        return gen_best_config(mat1, mat2)
+    else:
+        raise NotImplementedError("Function gen_best_config is not yet implemented")

outputs/audit_venv/lib/python3.11/site-packages/torch/_inductor/rocm_multiarch_utils.py

@@ -0,0 +1,310 @@
+"""
+ROCm Multi-Architecture Support Utilities
+Compile LLVM IR to multi-arch bundles that HIP can load automatically.
+"""
+
+import logging
+import os
+import re
+import subprocess
+
+import torch
+from torch.utils.cpp_extension import _join_rocm_home, ROCM_HOME
+
+
+log = logging.getLogger(__name__)
+
+
+def get_rocm_compiler() -> str:
+    """
+    Get path to ROCm's clang compiler.
+    Uses PyTorch's ROCM_HOME detection.
+
+    Returns:
+        Path to clang compiler
+
+    Raises:
+        RuntimeError: If ROCm is not found
+    """
+    if ROCM_HOME is None:
+        raise RuntimeError(
+            "ROCm installation not found. "
+            "PyTorch was not built with ROCm support or ROCM_HOME is not set."
+        )
+
+    # ROCm's clang is at <ROCM_HOME>/llvm/bin/clang
+    clang_path = _join_rocm_home("llvm", "bin", "clang")
+
+    if not os.path.exists(clang_path):
+        raise RuntimeError(
+            f"ROCm clang not found at {clang_path}. ROCM_HOME is set to {ROCM_HOME}"
+        )
+
+    return clang_path
+
+
+def get_rocm_bundler() -> str:
+    """
+    Get path to clang-offload-bundler.
+    Uses PyTorch's ROCM_HOME detection.
+
+    Returns:
+        Path to bundler
+
+    Raises:
+        RuntimeError: If bundler is not found
+    """
+    if ROCM_HOME is None:
+        raise RuntimeError(
+            "ROCm installation not found. "
+            "PyTorch was not built with ROCm support or ROCM_HOME is not set."
+        )
+
+    # Bundler is at <ROCM_HOME>/llvm/bin/clang-offload-bundler
+    bundler_path = _join_rocm_home("llvm", "bin", "clang-offload-bundler")
+
+    if not os.path.exists(bundler_path):
+        raise RuntimeError(
+            f"clang-offload-bundler not found at {bundler_path}. "
+            f"ROCM_HOME is set to {ROCM_HOME}"
+        )
+
+    return bundler_path
+
+
+def get_rocm_target_archs() -> list[str]:
+    env_archs = os.environ.get("PYTORCH_ROCM_ARCH", "").strip()
+    if env_archs:
+        archs = [arch.strip() for arch in env_archs.replace(";", ",").split(",")]
+        archs = [arch for arch in archs if arch]
+        if archs:
+            # Ensure current device arch is included
+            if torch.cuda.is_available():
+                for dev_idx in range(torch.cuda.device_count()):
+                    current_arch = torch.cuda.get_device_properties(
+                        dev_idx
+                    ).gcnArchName.split(":")[0]
+                    if current_arch not in archs:
+                        archs.append(current_arch)
+            return archs
+
+    try:
+        from torch._inductor import config
+
+        if hasattr(config, "rocm") and hasattr(config.rocm, "target_archs"):
+            archs = config.rocm.target_archs
+            if archs:
+                return archs
+
+    except Exception:
+        pass
+
+    return torch.cuda.get_arch_list()
+
+
+def _sanitize_llvm_ir_for_rocm(llvm_ir_path: str) -> str:
+    """
+    Sanitize LLVM IR to be compatible with ROCm's clang.
+
+    Triton's LLVM (upstream) may emit attributes and metadata that ROCm's
+    older clang does not yet support. Only strips attributes confirmed to
+    cause parse errors — preserves all others to maintain correct codegen.
+
+    Currently strips:
+        - nocreateundeforpoison: function attribute (upstream LLVM, not in ROCm)
+        - dwarfAddressSpace: debug metadata field (upstream LLVM, not in ROCm)
+
+    Returns:
+        Path to sanitized .ll file, or original path if no changes needed.
+    """
+    with open(llvm_ir_path) as f:
+        content = f.read()
+
+    sanitized = content
+    sanitized = re.sub(r"\bnocreateundeforpoison\b\s*", "", sanitized)
+    sanitized = re.sub(r",\s*dwarfAddressSpace:\s*\d+", "", sanitized)
+
+    if sanitized == content:
+        return llvm_ir_path
+
+    sanitized_path = llvm_ir_path + ".sanitized.ll"
+    with open(sanitized_path, "w") as f:
+        f.write(sanitized)
+
+    log.debug(
+        "Sanitized LLVM IR for ROCm clang compatibility: %s -> %s",
+        llvm_ir_path,
+        sanitized_path,
+    )
+    return sanitized_path
+
+
+def compile_llvm_ir_to_code_object(
+    llvm_ir_path: str, output_path: str, target_arch: str
+) -> bool:
+    """
+    Compile unbundled LLVM IR to a single-arch code object.
+
+    Args:
+        llvm_ir_path: Path to .ll file
+        output_path: Where to write .hsaco file
+        target_arch: Target architecture (e.g., 'gfx90a')
+
+    Returns:
+        True if successful
+    """
+    if not os.path.exists(llvm_ir_path):
+        return False
+
+    os.makedirs(os.path.dirname(output_path), exist_ok=True)
+
+    try:
+        clang = get_rocm_compiler()
+    except RuntimeError:
+        return False
+
+    # Sanitize LLVM IR to remove attributes unsupported by ROCm's clang
+    llvm_ir_path = _sanitize_llvm_ir_for_rocm(llvm_ir_path)
+
+    # Using clang and not hipcc since we are not compiling source code
+    # Instead we use the LLVM IR (.ll) provided by triton
+    cmd = [
+        clang,
+        "-target",
+        "amdgcn-amd-amdhsa",
+        f"-mcpu={target_arch}",
+        llvm_ir_path,
+        "-o",
+        output_path,
+    ]
+
+    try:
+        subprocess.run(cmd, capture_output=True, text=True, check=True)
+
+        if not os.path.exists(output_path):
+            return False
+
+        return True
+
+    except subprocess.CalledProcessError:
+        return False
+
+
+def create_multiarch_bundle(code_objects: dict, output_bundle_path: str) -> bool:
+    """
+    Bundle multiple architecture code objects into a single multi-arch bundle.
+
+    Uses clang-offload-bundler to create a fat binary that HIP runtime can load.
+    The runtime automatically selects the correct architecture at load time.
+
+    Args:
+        code_objects: Dict mapping architecture to code object path
+        output_bundle_path: Path for output bundle
+
+    Returns:
+        True if successful
+    """
+    if not code_objects:
+        return False
+
+    os.makedirs(os.path.dirname(output_bundle_path), exist_ok=True)
+
+    try:
+        bundler = get_rocm_bundler()
+    except RuntimeError:
+        return False
+
+    # Build targets and inputs lists for clang-offload-bundler
+    targets = ["host-x86_64-unknown-linux-gnu"]
+
+    # We include a dummy host entry to satisfy the bundler format
+    inputs = ["/dev/null"]
+
+    for arch, path in sorted(code_objects.items()):
+        if not os.path.exists(path):
+            continue
+        # hipv4 = HIP version 4 code object format
+        # amdgcn-amd-amdhsa = target triple for ROCm/HSA runtime
+        # arch = specific GPU (gfx90a, gfx942, etc.)
+        targets.append(f"hipv4-amdgcn-amd-amdhsa--{arch}")
+        inputs.append(path)
+
+    if len(inputs) == 1:  # Only host, no device code
+        return False
+
+    cmd = [
+        bundler,
+        "--type=o",
+        # CRITICAL: HIP runtime expects 4096-byte alignment for loading bundles
+        # Without this, hipModuleLoadData gives segmentation fault
+        "-bundle-align=4096",  # CRITICAL: Required by HIP runtime!
+        f"--targets={','.join(targets)}",
+    ]
+
+    for input_file in inputs:
+        cmd.append(f"--input={input_file}")
+
+    cmd.append(f"--output={output_bundle_path}")
+
+    try:
+        subprocess.run(cmd, capture_output=True, text=True, check=True)
+
+        if not os.path.exists(output_bundle_path):
+            return False
+
+        return True
+
+    except subprocess.CalledProcessError:
+        return False
+
+
+def compile_multiarch_bundle_from_llvm_ir(
+    llvm_ir_path: str, output_bundle_path: str, target_archs: list[str] | None = None
+) -> bool:
+    """
+    Complete workflow: LLVM IR → multiple code objects → bundle.
+
+    This is the main entry point for multi-arch compilation.
+
+    Args:
+        llvm_ir_path: Path to .ll file
+        output_bundle_path: Where to write bundle
+        target_archs: Optional list of architectures
+
+    Returns:
+        True if successful
+    """
+    if target_archs is None:
+        # Get architectures from environment variable or config
+        target_archs = get_rocm_target_archs()
+
+    # Step 1: Compile LLVM IR to code object for each architecture
+    code_objects = {}
+    temp_dir = os.path.dirname(output_bundle_path)
+    kernel_name = os.path.splitext(os.path.basename(llvm_ir_path))[0]
+
+    for arch in target_archs:
+        # Create temporary single-architecture code object
+        # Format: kernel_name_gfx90a.co, kernel_name_gfx942.co, etc.
+        co_path = os.path.join(temp_dir, f"{kernel_name}_{arch}.co")
+
+        # Compile with clang backend: LLVM IR → GPU machine code
+        if compile_llvm_ir_to_code_object(llvm_ir_path, co_path, arch):
+            code_objects[arch] = co_path
+
+    if not code_objects:
+        return False
+
+    # Step 2: Bundle all code objects together
+    # Uses clang-offload-bundler to create fat binary
+    success = create_multiarch_bundle(code_objects, output_bundle_path)
+
+    # Step 3: Clean up temporary single-arch code objects
+    # The bundle contains all the code, so intermediates are no longer needed
+    for co_path in code_objects.values():
+        try:
+            os.remove(co_path)
+        except Exception:
+            pass
+
+    return success

outputs/audit_venv/lib/python3.11/site-packages/torch/_inductor/runtime/autotune_cache.py

@@ -0,0 +1,667 @@
+"""
+PyTorch Inductor Autotuning Cache System
+
+This module implements a caching system for autotuning configurations in PyTorch's Inductor compiler.
+It provides mechanisms to store and retrieve optimal kernel configurations both locally and remotely,
+which significantly speeds up compilation by reusing previously discovered optimal parameters.
+
+The caching system includes:
+- Local filesystem caching for individual machine reuse
+- Remote caching for sharing optimizations across machines
+- Bundled caching to efficiently store multiple related configurations
+- Cache invalidation based on PyTorch versions and backend changes
+- Serialization/deserialization support for worker processes
+
+Key components:
+- AutotuneCache: Main class for managing cache access and storage
+- AutotuneCacheBundler: Bundles multiple cache entries for efficient storage
+- LocalAutotuneCache: Handles filesystem-based caching
+- _LocalAutotuneCacheBackend: Low-level file operations for cache storage
+- AutotuneCacheArtifact: Integration with PyTorch's artifact system
+
+This caching system is critical for performance as it eliminates the need to re-run
+expensive autotuning operations when the same kernels are compiled multiple times.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import hashlib
+import logging
+import os
+import os.path
+import re
+from typing import Any, TYPE_CHECKING
+from typing_extensions import override
+
+import torch
+from torch._inductor.runtime.runtime_utils import cache_dir
+from torch.compiler._cache import (
+    CacheArtifact,
+    CacheArtifactFactory,
+    CacheArtifactManager,
+)
+from torch.utils._triton import has_triton
+
+from ..remote_cache import (
+    create_cache,
+    JsonDataTy,
+    RemoteCache,
+    RemoteCacheBackend,
+    RemoteCacheJsonSerde,
+)
+from .triton_compat import Config, HAS_WARP_SPEC
+
+
+if TYPE_CHECKING:
+    from ..remote_cache import Sample
+
+log = logging.getLogger(__name__)
+
+
+_InductorMetaTy = dict[str, object]
+
+
+def inductor_meta_from_config() -> _InductorMetaTy:
+    from torch._inductor import config
+
+    backend_hash = None
+    if has_triton():
+        try:
+            backend_hash = torch.utils._triton.triton_hash_with_backend()
+        except RuntimeError:
+            # This can get the error:
+            #   RuntimeError: 0 active drivers ([]). There should only be one.
+            pass
+
+    is_hip = None
+    if torch.version.hip is not None:
+        is_hip = True
+
+    return {
+        "autotune_local_cache": config.autotune_local_cache,
+        "autotune_remote_cache": config.autotune_remote_cache,
+        "backend_hash": backend_hash,
+        "bundled_autotune_remote_cache": config.bundled_autotune_remote_cache,
+        "coordinate_descent_tuning": config.coordinate_descent_tuning,
+        "is_fbcode": config.is_fbcode(),
+        "is_hip": is_hip,
+    }
+
+
+@CacheArtifactFactory.register
+class AutotuneCacheArtifact(CacheArtifact):
+    @override
+    def populate_cache(self) -> None:
+        autotune_cache = _LocalAutotuneCacheBackend()
+        key = os.path.join(cache_dir(), self.key)
+        autotune_cache._put(key, self.content)
+
+    @override
+    @staticmethod
+    def type() -> str:
+        return "autotune"
+
+    @override
+    @staticmethod
+    def encode(content: JsonDataTy) -> bytes:
+        assert not isinstance(content, bytes)
+        serde = RemoteCacheJsonSerde()
+        content_bytes = serde.encode(content)
+        assert isinstance(content_bytes, bytes)
+        return content_bytes
+
+
+@dataclasses.dataclass
+class AutotuneCache:
+    configs_hash: str
+    local_cache: tuple[RemoteCache[JsonDataTy], str] | None = None
+    remote_cache: tuple[RemoteCache[JsonDataTy], str] | None = None
+
+    # Create a AutotuneCache. Returns None if none of the caches can be used.
+    @staticmethod
+    def create(
+        inductor_meta: _InductorMetaTy, filename: str, configs_hash: str
+    ) -> AutotuneCache | None:
+        cache = AutotuneCache(configs_hash)
+        key = AutotuneCache._prepare_key(filename)
+
+        cache._setup_local_cache(inductor_meta, os.path.dirname(filename), key)
+        cache._setup_remote_autotune_cache(inductor_meta, key)
+        if cache.local_cache or cache.remote_cache:
+            return cache
+        else:
+            return None
+
+    @staticmethod
+    def _prepare_key(filename: str) -> str:
+        from torch.compiler import config as cconfig
+
+        # base of filename is already sha256 hash the source contents
+        key = f"{os.path.basename(filename)}:{cconfig.cache_key_tag}"
+        return hashlib.sha256(key.encode("utf-8")).hexdigest()
+
+    # Read the best config options from the most local cache and return it.
+    def _read(self) -> dict[str, JsonDataTy] | None:
+        if local_cache := self.local_cache:
+            cache, key = local_cache
+            if best_config := cache.get(key):
+                if isinstance(best_config, dict):
+                    return best_config
+
+        if remote_cache := self.remote_cache:
+            cache, key = remote_cache
+            if best_config := cache.get(key):
+                if isinstance(best_config, dict):
+                    return best_config
+
+        return None
+
+    # Read the best config options from the most local cache and figure out
+    # which `configs` represents that option.
+    def read_best(
+        self, inductor_meta: _InductorMetaTy, configs: list[Config]
+    ) -> Config | None:
+        if best := self._read():
+            return _load_cached_autotuning(
+                best, self.configs_hash, configs, inductor_meta
+            )
+        return None
+
+    # Set up local filesystem caching information
+    def _setup_local_cache(
+        self, inductor_meta: _InductorMetaTy, dirname: str, cache_key: str
+    ) -> None:
+        if not inductor_meta.get("autotune_local_cache", True):
+            return
+
+        from ..codecache import torch_key
+
+        """
+        [Note: torch_key in autotune cache key]
+        Include torch_key() in the cache key so that different versions
+        of torch result in cache invalidation. This is important in case
+        of changes to the best_config format or other code changes that
+        are not backward compatible w.r.t. the cache.
+        """
+        hasher = hashlib.sha256()
+        hasher.update(cache_key.encode("utf-8"))
+        hasher.update(torch_key())
+        updated_cache_key = hasher.hexdigest()
+
+        cache_filename = f"{dirname}/{updated_cache_key}.best_config"
+        local_cache = LocalAutotuneCache()
+        self.local_cache = (local_cache, cache_filename)
+
+    # Set up remote caching information
+    def _setup_remote_autotune_cache(
+        self, inductor_meta: _InductorMetaTy, cache_key: str
+    ) -> None:
+        if not _should_use_remote_autotune_cache(inductor_meta):
+            return
+
+        if (backend_hash := inductor_meta.get("backend_hash", None)) is None:
+            log.debug(
+                "backend_hash is not passed on the inductor_meta, unable to use autotune remote cache"
+            )
+            return
+        assert isinstance(backend_hash, str)
+
+        from ..codecache import torch_key
+
+        is_fbcode = bool(inductor_meta.get("is_fbcode", False))
+
+        salt = "autotune-best-config-v2"
+        # re: torch_key - see [Note: torch_key in autotune cache key]
+        key = torch_key().hex() + backend_hash + self.configs_hash + salt
+        key = hashlib.sha256(key.encode("utf-8")).hexdigest()
+
+        remote_cache = create_cache(
+            key,
+            is_fbcode,
+            "FbRemoteAutotuneCache",
+            "RemoteAutotuneCache",
+        )
+        if not remote_cache:
+            return
+
+        # Save the args passed to create_cache
+        # in case AutotuneCache needs to be pickled
+        self.remote_cache_full_key = key
+        self.is_fbcode = is_fbcode
+        self.remote_cache = (remote_cache, cache_key)
+
+    # The AutotuneCache may be serialized/deserialized if we're using
+    # AsyncCompile worker processes to run triton compilation.
+    # This is because AutotuneCache instances are created on the worker
+    # process, but we need to run AutotuneCache.save on the parent process
+    # when actually doing autotuning.
+    def __getstate__(self) -> dict[str, Any]:
+        # The remote cache handles themselves may not be serializable
+        # So clear it and reconstruct it on setstate
+        remote_cache = getattr(self, "remote_cache", None)
+        return {
+            **self.__dict__,
+            # Save the cache_key portion
+            "remote_cache": remote_cache and remote_cache[1],
+        }
+
+    def __setstate__(self, state: dict[str, Any]) -> None:
+        # Reconstruct the remote cache on the parent class
+        self.__dict__.update(state)
+        if self.remote_cache is not None:
+            assert isinstance(self.remote_cache, str)
+            assert hasattr(self, "remote_cache_full_key")
+            assert hasattr(self, "is_fbcode")
+            cache_key = self.remote_cache
+            remote_cache = create_cache(
+                self.remote_cache_full_key,
+                self.is_fbcode,
+                "FbRemoteAutotuneCache",
+                "RemoteAutotuneCache",
+            )
+            if remote_cache is not None:
+                self.remote_cache = (remote_cache, cache_key)
+            else:
+                log.warning("Warning, failed to recreate remote cache after pickling")
+                self.remote_cache = None
+
+    # Save the config in the caches
+    def save(
+        self,
+        config: Config,
+        time_taken_ns: int,
+        found_by_coordesc: bool = False,
+        triton_cache_hash: str | None = None,
+    ) -> None:
+        data = {
+            # pyrefly: ignore [missing-attribute]
+            **config.kwargs,
+            # pyrefly: ignore [missing-attribute]
+            "num_warps": config.num_warps,
+            # pyrefly: ignore [missing-attribute]
+            "num_stages": config.num_stages,
+            "configs_hash": self.configs_hash,
+            "found_by_coordesc": found_by_coordesc,
+            "time_taken_ms": time_taken_ns // 1000000,  # Convert from NS to MS
+            "triton_cache_hash": triton_cache_hash,
+        }
+        # Save extra_options if present on the config. This allows third-party
+        # backends to store custom tuned options alongside the standard config.
+        if extra_options := getattr(config, "extra_options", None):
+            data["extra_options"] = extra_options
+        if HAS_WARP_SPEC:
+            data.update(
+                {
+                    "num_consumer_groups": getattr(config, "num_consumer_groups", 0),
+                    "num_buffers_warp_spec": getattr(
+                        config, "num_buffers_warp_spec", 0
+                    ),
+                }
+            )
+
+        if local_cache := self.local_cache:
+            cache, key = local_cache
+            # pyrefly: ignore [bad-argument-type]
+            cache.put(key, data)
+            # pyrefly: ignore [bad-argument-type]
+            AutotuneCacheBundler.put(key, data)
+            autotune_artifact_key = os.path.join(*key.split(os.sep)[-2:])
+            CacheArtifactManager.record_artifact(
+                AutotuneCacheArtifact.type(), autotune_artifact_key, data
+            )
+
+            if log.isEnabledFor(logging.DEBUG):
+                type_str = "coordesc" if found_by_coordesc else "heuristic"
+                log.debug("Save %s tuning result to %s", type_str, key)
+
+        if remote_cache := self.remote_cache:
+            cache, key = remote_cache
+            # pyrefly: ignore [bad-argument-type]
+            cache.put(key, data)
+
+
+class _AutotuneCacheBundlerImpl:
+    """
+    Caches a set of LocalAutotuneCacheBackend entries together in a single
+    cache.
+    """
+
+    _key: str
+    _cache: RemoteCache[JsonDataTy]
+
+    # All known entries from LocalAutotuneCache.put()
+    _entries: dict[str, JsonDataTy]
+
+    def end_compile(self) -> None:
+        # TODO: Do we need to compute time_taken_ms and encode that somehow?
+        if self._entries:
+            self._cache.put(self._key, self._entries)
+
+    def put(self, basename: str, data: JsonDataTy) -> None:
+        # Do we need to worry about duplicates? We only have a single local fs
+        # entry - so probably not.
+        self._entries[basename] = data
+
+    def __init__(self, key: str, cache: RemoteCache[JsonDataTy]) -> None:
+        self._key = key
+        self._cache = cache
+        self._entries = {}
+
+    def sync(self) -> None:
+        # We don't currently use this - but we could async load starting at
+        # `begin_compile` and wait for the load to be finished here.
+        pass
+
+    @classmethod
+    def _should_use_bundled_autotune_remote_cache(
+        cls, inductor_meta: _InductorMetaTy
+    ) -> bool:
+        # The bundled autotune cache is only available if you've also got local
+        # caching enabled (because we feed the bundled data to the local cache).
+        if not inductor_meta.get("autotune_local_cache", True):
+            return False
+
+        # Check if the we're enabled via config
+        if (
+            bundled_autotune_remote_cache := inductor_meta.get(
+                "bundled_autotune_remote_cache"
+            )
+        ) is not None:
+            return bool(bundled_autotune_remote_cache)
+
+        if not cls._get_is_fbcode(inductor_meta):
+            return False
+        if torch._utils_internal.is_fb_unit_test():
+            return False
+        if inductor_meta.get("is_hip"):
+            return False
+
+        try:
+            from torch._inductor.fb.remote_cache import REMOTE_CACHE_VERSION
+        except ModuleNotFoundError:
+            return False
+
+        jk = torch._utils_internal.justknobs_getval_int(
+            "pytorch/remote_cache:bundled_autotune_remote_cache_version"
+        )
+        return REMOTE_CACHE_VERSION >= jk
+
+    def _load_cache(self) -> bool:
+        from torch._inductor import codecache
+
+        # The single key is defined on construction of the cache.
+        entries = self._cache.get(self._key)
+        if entries is None or not isinstance(entries, dict):
+            # We couldn't load the cache - so mark _entries as non-None so we
+            # store local cache values.
+            return False
+
+        # Go through the entries we got from the cache and save them locally.
+        time_saved_ns = 0
+        for basename, data in entries.items():
+            # Reconstruct the final filename (see put())
+            root, ext = _splitext_nodot(basename)
+            _, _, filename = codecache.get_path(root, ext)
+            if isinstance(data, dict) and (tsns := data.get("time_saved_ns")):
+                time_saved_ns += int(tsns)  # type: ignore[arg-type]
+            local_cache = LocalAutotuneCache()
+            local_cache.put(filename, data)
+
+        codecache.add_ephemeral_timeout_increase_for_distributed(time_saved_ns)
+
+        return True
+
+    @staticmethod
+    def _get_is_fbcode(inductor_meta: _InductorMetaTy) -> bool:
+        return bool(inductor_meta.get("is_fbcode", False))
+
+    @staticmethod
+    def _get_backend_hash(inductor_meta: _InductorMetaTy) -> str:
+        backend_hash = inductor_meta["backend_hash"]
+        assert isinstance(backend_hash, str)
+        return backend_hash
+
+
+class AutotuneCacheBundler:
+    _bundler: _AutotuneCacheBundlerImpl | None = None
+
+    def __init__(self) -> None:
+        pass
+
+    # Call this before we start any autotune computation for an inductor python
+    # file. On a cache hit it copies the individual results into the local
+    # autotune caches.
+    @classmethod
+    def begin_compile(
+        cls,
+        inductor_meta: _InductorMetaTy,
+        *,
+        code: str | None = None,
+        code_hash: str | None = None,
+    ) -> None:
+        assert cls._bundler is None
+
+        if code is not None:
+            assert code_hash is None, "Cannot specify both code and code_hash"
+            code_hash = _comment_stripped_hash(code)
+        assert code_hash is not None
+
+        if not _AutotuneCacheBundlerImpl._should_use_bundled_autotune_remote_cache(
+            inductor_meta
+        ):
+            return
+
+        cache = create_cache(
+            "bundled-autotune-v1",
+            _AutotuneCacheBundlerImpl._get_is_fbcode(inductor_meta),
+            "FbRemoteBundledAutotuneCache",
+            "RemoteBundledAutotuneCache",
+        )
+        if not cache:
+            return
+
+        # We're starting a compilation phase. We have a cache key for the code
+        # we're compiling. We'll get the individual autotune bundles later (via
+        # self.put()). For now create the AutotuneCacheBundler and try to load
+        # from the cache.
+
+        salt = "bundled-autotune-best-configs-v1"
+        backend_hash = _AutotuneCacheBundlerImpl._get_backend_hash(inductor_meta)
+        # TODO: The autotune cache includes configs_hash in the key. The problem
+        # is that the configs_hash includes info from the individual pointwise()
+        # calls (size_hints, for example) which we can't know yet. I *think*
+        # that info is basically present in the `code_hash` (since it's a
+        # parameter to the pointwise decorator) - but is there other info we
+        # need to include from inductor_meta?
+        key = code_hash + backend_hash + salt
+        key = hashlib.sha256(key.encode("utf-8")).hexdigest()
+
+        bundler = _AutotuneCacheBundlerImpl(key, cache)
+        if not bundler._load_cache():
+            # We couldn't load from the cache - so save the data so we can store
+            # the saved autotunes.
+            cls._bundler = bundler
+
+        # If we get a cache hit don't bother saving any of the individual
+        # autotune results.
+
+    # Call this after all individual autotune results are finished for a
+    # inductor python file. If we gathered any individual results then we bundle
+    # those and put it into the cache.
+    @classmethod
+    def end_compile(cls) -> None:
+        if bundler := cls._bundler:
+            cls._bundler = None
+            bundler.end_compile()
+
+    @classmethod
+    def sync(cls) -> None:
+        if bundler := cls._bundler:
+            bundler.sync()
+
+    @classmethod
+    def put(cls, filename: str, data: JsonDataTy) -> None:
+        if bundler := cls._bundler:
+            # The filename comes in as something like
+            # "/tmp/tmp{random}/{aa}/{basename}.py" (where aa is
+            # basename[1:3]). Strip it down and make sure that it looks like a path
+            # we could reconstruct (because it's possible for the caller to
+            # customize the path).
+            basename = os.path.basename(filename)
+
+            # TODO: check cache_dir() vs filename, then strip dirname
+            bundler.put(basename, data)
+
+
+# Remove the comments from the code (which include things like run ids and file
+# paths) and then hash the result.
+def _comment_stripped_hash(code: str) -> str:
+    code = re.sub(r"#.*$", "", code, count=0, flags=re.MULTILINE)
+    return torch._inductor.codecache.code_hash(code)
+
+
+def _should_use_remote_autotune_cache(inductor_meta: _InductorMetaTy) -> bool:
+    if (config := inductor_meta.get("autotune_remote_cache")) is not None:
+        return bool(config)
+    if not inductor_meta.get("is_fbcode"):
+        return False
+    if torch._utils_internal.is_fb_unit_test():
+        return False
+    if inductor_meta.get("is_hip"):
+        return False
+
+    try:
+        from torch._inductor.fb.remote_cache import REMOTE_CACHE_VERSION
+    except ModuleNotFoundError:
+        return False
+
+    return REMOTE_CACHE_VERSION >= torch._utils_internal.justknobs_getval_int(
+        "pytorch/remote_cache:autotune_memcache_version"
+    )
+
+
+def _load_cached_autotuning(
+    best_config: dict[str, JsonDataTy],
+    configs_hash: str,
+    configs: list[Config],
+    inductor_meta: _InductorMetaTy,
+) -> Config | None:
+    if best_config is None:
+        return None
+    if best_config.pop("configs_hash", None) != configs_hash:
+        return None
+
+    # Remove time taken for comparison
+    best_config.pop("time_taken_ms", None)
+
+    best_config.pop("triton_cache_hash", None)
+
+    # Extract extra_options if present. This allows third-party backends
+    # to restore custom tuned options from the cache.
+    extra_options = best_config.pop("extra_options", None)
+
+    if inductor_meta.get("coordinate_descent_tuning") and best_config.pop(
+        "found_by_coordesc", False
+    ):
+        num_warps = best_config.pop("num_warps")
+        num_stages = best_config.pop("num_stages")
+
+        # Extract common arguments
+        config_args = {
+            "num_warps": num_warps,
+            "num_stages": num_stages,
+        }
+
+        if HAS_WARP_SPEC:
+            config_args.update(
+                {
+                    "num_consumer_groups": best_config.pop("num_consumer_groups", 0),
+                    "num_buffers_warp_spec": best_config.pop(
+                        "num_buffers_warp_spec", 0
+                    ),
+                }
+            )
+
+        # Create the triton_config with the appropriate arguments
+        # pyrefly: ignore [bad-argument-count, unexpected-keyword]
+        triton_config = Config(best_config, **config_args)
+        # pyrefly: ignore [missing-attribute]
+        triton_config.found_by_coordesc = True
+        # Restore extra_options (may be None if not used by backend)
+        # pyrefly: ignore [missing-attribute]
+        triton_config.extra_options = extra_options
+        return triton_config
+
+    matching_configs = [
+        cfg
+        for cfg in configs
+        # pyrefly: ignore [missing-attribute]
+        if all(val == best_config.get(key) for key, val in cfg.kwargs.items())
+        # pyrefly: ignore [missing-attribute]
+        and cfg.num_warps == best_config.get("num_warps")
+        # pyrefly: ignore [missing-attribute]
+        and cfg.num_stages == best_config.get("num_stages")
+    ]
+    if len(matching_configs) != 1:
+        return None
+
+    matched_config = matching_configs[0]
+    # Restore extra_options (may be None if not used by backend)
+    # pyrefly: ignore [missing-attribute]
+    matched_config.extra_options = extra_options
+    return matched_config
+
+
+class _LocalAutotuneCacheBackend(RemoteCacheBackend[bytes]):
+    @override
+    def _get(self, key: str) -> bytes | None:
+        try:
+            with open(key, "rb") as fd:
+                return fd.read()
+        except FileNotFoundError:
+            return None
+
+    @override
+    def _put(self, key: str, data: bytes) -> None:
+        os.makedirs(os.path.dirname(key), exist_ok=True)
+        from torch._inductor import codecache
+
+        codecache.write_atomic(key, data)
+
+
+class LocalAutotuneCache(RemoteCache[JsonDataTy]):
+    def __init__(self) -> None:
+        backend = _LocalAutotuneCacheBackend()
+        serde = RemoteCacheJsonSerde()
+        super().__init__(backend, serde)
+
+    @override
+    def _get(self, key: str, sample: Sample | None) -> JsonDataTy | None:
+        AutotuneCacheBundler.sync()
+        result = super()._get(key, sample)
+        if result is not None:
+            assert isinstance(result, dict)
+            # What? Why are we doing a put() here? Imagine we have a new model
+            # that reuses some existing kernels that have already been
+            # compiled. If we didn't do a `put` here (on cache hit) then the new
+            # model would only bundle *newly* compiled kernels, not existing
+            # kernels that were already compiled and cached.
+            AutotuneCacheBundler.put(key, result)
+            autotune_artifact_key = os.path.join(*key.split(os.sep)[-2:])
+            CacheArtifactManager.record_artifact(
+                AutotuneCacheArtifact.type(), autotune_artifact_key, result
+            )
+        return result
+
+    @override
+    def _put(self, key: str, value: JsonDataTy, sample: Sample | None) -> None:
+        AutotuneCacheBundler.put(key, value)
+        super()._put(key, value, sample)
+
+
+def _splitext_nodot(basename: str) -> tuple[str, str]:
+    root, ext = os.path.splitext(basename)
+    if ext:
+        ext = ext[1:]
+    return root, ext

outputs/audit_venv/lib/python3.11/site-packages/torch/_inductor/runtime/benchmarking.py

@@ -0,0 +1,549 @@
+import functools
+import inspect
+import time
+from collections.abc import Callable
+from functools import cached_property, wraps
+from itertools import chain
+from statistics import median
+from typing import Any, Concatenate
+from typing_extensions import ParamSpec, Self, TypeVar
+
+import torch
+import torch._inductor.config as inductor_config
+import torch.utils._pytree as pytree
+from torch._dynamo.utils import counters
+from torch.utils._debug_mode import DebugMode
+
+
+logger = torch._logging.getArtifactLogger(__name__, "benchmarking")
+use_experimental_benchmarker = (
+    inductor_config.use_experimental_benchmarker and torch.cuda.is_available()
+)
+
+
+MILLISECONDS_PER_SECOND = 1000
+
+P = ParamSpec("P")
+T = TypeVar("T")
+
+
+# Device-type → benchmarking function registry.
+# Keys must match torch.device.type (e.g., "cpu", "cuda", "mps", "xpu", ...).
+# Values are callables with signature:
+#   fn(self: Benchmarker, _callable: Callable[..., Any], *, warmup: int, rep: int, **kwargs) -> Any
+_BENCHMARK_DISPATCH: dict[str, Callable[..., Any]] = {}
+
+
+def register_benchmarker(
+    device_type: str,
+    fn: Callable[..., Any],
+    *,
+    override: bool = False,
+) -> None:
+    """
+    Register a device-type specific benchmarker.
+
+    Args:
+        device_type: torch.device.type string (e.g., "cuda", "cpu", "mps", "xpu").
+        fn: callable(self, _callable, *, warmup, rep, **kwargs) -> Any
+        override: allow overriding an existing registration.
+    """
+    if not isinstance(device_type, str) or not device_type:
+        raise ValueError(
+            "device_type must be a non-empty string matching torch.device.type"
+        )
+    if not callable(fn):
+        raise TypeError("fn must be callable")
+    if not override and device_type in _BENCHMARK_DISPATCH:
+        raise ValueError(
+            f"Benchmarker for device_type '{device_type}' already registered"
+        )
+    _BENCHMARK_DISPATCH[device_type] = fn
+
+
+def may_distort_benchmarking_result(fn: Callable[..., Any]) -> Callable[..., Any]:
+    from torch._inductor import config
+
+    if config.test_configs.distort_benchmarking_result == "":
+        return fn
+
+    def distort(
+        ms: list[float] | tuple[float, ...] | float,
+    ) -> list[float] | tuple[float, ...] | float:
+        if isinstance(ms, (list, tuple)):
+            return type(ms)(distort(val) for val in ms)  # type: ignore[misc]
+
+        distort_method = config.test_configs.distort_benchmarking_result
+        assert isinstance(ms, float)
+        if distort_method == "inverse":
+            return 1.0 / ms if ms else 0.0
+        elif distort_method == "random":
+            import random
+
+            return random.random()
+        else:
+            raise RuntimeError(f"Unrecognized distort method {distort_method}")
+
+    @functools.wraps(fn)
+    def wrapper(
+        *args: list[Any], **kwargs: dict[str, Any]
+    ) -> list[float] | tuple[float, ...] | float:
+        ms = fn(*args, **kwargs)
+
+        return distort(ms)
+
+    return wrapper
+
+
+def may_ban_benchmarking() -> None:
+    if torch._inductor.config.deterministic:
+        raise RuntimeError("""In the deterministic mode of Inductor, we will avoid those
+        benchmarkings that would cause non deterministic results. Only benchmarkings in the vetted
+        scenarios are allowed. Example include autotuning for triton configs of pointwise kernels.
+
+        When you see this exception, you can do one of the following two things:
+        1. if the benchmarking you are doing does not introduce any non-determinism, you can just
+        add is_vetted_benchmarking=True to you benchmark_gpu call. That would solve the issue.
+
+        2. if the benchmarking you are doing indeed introduces non-determinism, you'll need to disable
+        such feature in deterministic mode or find an alternative implementation that is deterministic.
+        """)
+
+
+def time_and_count(
+    fn: Callable[Concatenate[Any, P], T],
+) -> Callable[Concatenate[Any, P], T]:
+    """
+    Wraps `fn` to increment the appropriate dynamo counters. It is expected that `fn`
+    is a method of `Benchmarker` or one of its subclasses; typing limitations prevent
+    us from declaring this directly.
+
+    NOTE: If you're tempted to add a dynamo_timed call here, this function can be
+    called enough that the dynamo_timed overhead is not negligible.
+    """
+
+    @wraps(fn)
+    def wrapper(self: Any, *args: P.args, **kwargs: P.kwargs) -> T:
+        fn_qual_name = f"{self.__class__.__name__}.{fn.__name__}"
+        counters["inductor"][f"benchmarking.{fn_qual_name}"] += 1
+        return fn(self, *args, **kwargs)
+
+    return wrapper
+
+
+class Benchmarker:
+    """
+    A device-agnostic benchmarking utility for measuring the runtime of
+    inductor generated callables.
+    """
+
+    def __init__(self: Self) -> None:
+        pass
+
+    def infer_device(self, *fn_args: Any, **fn_kwargs: Any) -> torch.device:
+        inferred_device: torch.device | None = None
+        for arg_or_kwarg in chain(fn_args, fn_kwargs.values()):
+            # Some callables take nested structures as arguments so use the
+            # flattened form to find any tensors
+            for arg_or_kwarg_leaf in pytree.tree_leaves(arg_or_kwarg):
+                if not isinstance(arg_or_kwarg_leaf, torch.Tensor):
+                    continue
+                if inferred_device is None:
+                    inferred_device = arg_or_kwarg_leaf.device
+                elif arg_or_kwarg_leaf.device != inferred_device:
+                    raise ValueError(
+                        "Can't safely infer the device type of `fn` with multiple device types in `fn_args` and `fn_kwargs`!"
+                    )
+
+        if inferred_device is None:
+            raise ValueError(
+                "Can't safely infer the device type of `fn` with no device types"
+                " in `fn_args` or `fn_kwargs`. Use a direct benchmarking method instead e.g. "
+                "`Benchmarker.benchmark_cpu` or `Benchmarker.benchmark_gpu`."
+            )
+
+        return inferred_device
+
+    @time_and_count
+    def benchmark(
+        self: Self,
+        fn: Callable[..., Any],
+        fn_args: tuple[Any, ...] | None = None,
+        fn_kwargs: dict[str, Any] | None = None,
+        device: str | torch.device | None = None,
+        **kwargs: Any,
+    ) -> float:
+        """Benchmark `fn(*fn_args, *fn_kwargs)` and return the runtime, in milliseconds (the
+        actual runtime calculation is dictated by the benchmarking implementation, but may be
+        one of [mean, median, minimum, etc.]). Functions as a convenience wrapper around
+        device-specific implementations, like `benchmark_cpu` and `benchmark_gpu`. Raises
+        `ValueError(...)` if we can't safely infer the device type of `fn`; for example,
+        if multiple device types are found in `fn_args` and `fn_kwargs`, or if no device
+        types are found. To bypass device inference, provide the device to the `device`
+        parameter.
+
+        WARNING: if `fn` mutates `fn_args` or `fn_kwargs`, benchmarking may fail unexpectedly.
+        For example, if `fn` clears a mutable object, subsequent invocations of `fn` during
+        benchmarking will fail. In such cases, `fn` should handle cloning its arguments internally.
+        If device inference is required, `Benchmarker.infer_device` can be used prior to calling
+        this method without any arguments for `fn_args` and `fn_kwargs`.
+
+        Arguments:
+        - fn: The function to benchmark.
+        - fn_args: The function's arguments.
+        - fn_kwargs: The function's kwargs.
+
+        Keyword Arguments:
+        - device: Which device to use for benchmarking. If not provided the device will be attempted
+        to be inferred from `fn_args` and `fn_kwargs`.
+        - **kwargs: The benchmarking implementation's kwargs.
+
+        Returns:
+        - The runtime of `fn(*fn_args, **fn_kwargs)`, in milliseconds.
+        """
+        inferred_device: torch.device | None = None
+        if device is not None:
+            inferred_device = (
+                torch.device(device) if isinstance(device, str) else device
+            )
+        else:
+            if fn_args is None and fn_kwargs is None:
+                raise ValueError(
+                    "`fn_args` and `fn_kwargs` cannot both be None if `device` is not provided."
+                )
+
+            fn_args = fn_args or tuple()
+            fn_kwargs = fn_kwargs or {}
+            inferred_device = self.infer_device(*fn_args, **fn_kwargs)
+
+        assert isinstance(inferred_device, torch.device)
+
+        fn_args = fn_args or tuple()
+        fn_kwargs = fn_kwargs or {}
+
+        # No need to wrap if the callable takes no arguments
+        if len(fn_args) == 0 and len(fn_kwargs) == 0:
+            # Keep a true zero-arg callable type to satisfy type checkers.
+            def _callable() -> Any:
+                return fn()
+        else:
+            _args = fn_args
+            _kwargs = fn_kwargs
+
+            def _callable() -> Any:
+                return fn(*_args, **_kwargs)
+
+        warmup = kwargs.pop("warmup", inductor_config.inductor_default_autotune_warmup)
+        rep = kwargs.pop("rep", inductor_config.inductor_default_autotune_rep)
+
+        # Surfacing all kernels during autotuning is super noisy; filtering these out.
+        with DebugMode._benchmarking_inductor():
+            # First, try a registered device-specific benchmarker
+            benchmark_fn: Callable[..., Any] | None = _BENCHMARK_DISPATCH.get(
+                inferred_device.type
+            )
+            if benchmark_fn is not None:
+                return benchmark_fn(self, _callable, warmup=warmup, rep=rep, **kwargs)
+
+            # Backward-compatible default:
+            # - CPU  -> CPU benchmark path
+            # - else -> GPU benchmark path (legacy behavior retained for non-CPU)
+            if inferred_device == torch.device("cpu"):
+                return self.benchmark_cpu(_callable, warmup=warmup, rep=rep, **kwargs)
+            return self.benchmark_gpu(_callable, warmup=warmup, rep=rep, **kwargs)
+
+    @time_and_count
+    def benchmark_cpu(
+        self: Self, _callable: Callable[[], Any], warmup: int = 20, rep: int = 100
+    ) -> float:
+        """Benchmark the CPU callable, `_callable`, and return the median runtime,
+        in milliseconds.
+
+        Arguments:
+        - _callable: The CPU callable to benchmark.
+
+        Keyword Arguments:
+        - warmup: Optionally, the duration, in milliseconds, to run `_callable`
+        before benchmarking starts.
+        - rep: Optionally, the duration, in milliseconds, to run `_callable`
+        during benchmarking.
+
+        Returns:
+        - The median runtime of `_callable`, in milliseconds.
+        """
+
+        def run_for(ms: int) -> list[float]:
+            timings = []
+            run_start_t = time.perf_counter()
+            while True:
+                start_t = time.perf_counter()
+                _callable()
+                end_t = time.perf_counter()
+                timings.append((end_t - start_t) * MILLISECONDS_PER_SECOND)
+                if ((end_t - run_start_t) * MILLISECONDS_PER_SECOND) > ms:
+                    break
+            return timings
+
+        run_for(warmup)
+        return median(run_for(rep))
+
+    @time_and_count
+    def benchmark_gpu(self: Self, *args: Any, **kwargs: Any) -> float:
+        raise NotImplementedError
+
+    @time_and_count
+    def benchmark_gpu_with_cuda_graph(
+        self: Self,
+        _callable: Callable[[], Any],
+        **kwargs: Any,
+    ) -> float:
+        """Benchmark a GPU callable using CUDA graph capture and replay.
+
+        This captures the callable into a CUDA graph and benchmarks the graph replay,
+        which eliminates kernel launch overhead for fair comparison between different
+        implementations.
+        """
+        # Warmup
+        _callable()
+        torch.cuda.synchronize()
+
+        # Capture into CUDA graph
+        cuda_graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(cuda_graph, capture_error_mode="thread_local"):
+            _callable()
+        torch.cuda.synchronize()
+
+        return self.benchmark_gpu(cuda_graph.replay, **kwargs)
+
+
+# Make built-in defaults explicit via the registry
+def _default_cpu_bench(self, f, *, warmup, rep, **kw):
+    return self.benchmark_cpu(f, warmup=warmup, rep=rep, **kw)
+
+
+def _default_cuda_bench(self, f, *, warmup, rep, **kw):
+    return self.benchmark_gpu(f, warmup=warmup, rep=rep, **kw)
+
+
+def _default_xpu_bench(self, f, *, warmup, rep, **kw):
+    return self.benchmark_gpu(f, warmup=warmup, rep=rep, **kw)
+
+
+register_benchmarker("cpu", _default_cpu_bench, override=True)
+register_benchmarker("cuda", _default_cuda_bench, override=True)
+register_benchmarker("xpu", _default_xpu_bench, override=True)
+
+
+class TritonBenchmarker(Benchmarker):
+    @cached_property
+    def triton_do_bench(self: Self) -> Callable[..., Any]:
+        """Lazily import Triton's `do_bench`."""
+        try:
+            from triton.testing import do_bench
+        except ImportError as e:
+            raise NotImplementedError("requires Triton") from e
+        return do_bench
+
+    @may_distort_benchmarking_result
+    @time_and_count
+    # pyrefly: ignore [bad-override]
+    def benchmark_gpu(
+        self: Self,
+        _callable: Callable[[], Any],
+        is_vetted_benchmarking: bool = False,
+        **kwargs: Any,
+    ) -> float:
+        """Benchmark the GPU callable, `_callable`, and return the runtime, in milliseconds.
+
+        Arguments:
+        - _callable: The GPU callable to benchmark.
+
+        Keyword Arguments:
+        - quantiles: Optionally, a tuple of floats denoting the requested quantiles.
+        - return_mode: Optionally, the requested return mode. Currently, Triton's
+        `do_bench` supports min, max, mean, and median return modes.
+        - **kwargs: Additional kwargs passed to Triton's `do_bench`.
+
+        Returns:
+        - The runtime of `callable`, in milliseconds. If `kwargs["quantiles"]` is specified,
+        this is the first requested quantile. Else, if `kwargs["return_mode"]` is specified,
+        this is the requested return mode. Otherwise, this is the median.
+        """
+        if not is_vetted_benchmarking:
+            may_ban_benchmarking()
+
+        do_bench_params = inspect.signature(self.triton_do_bench).parameters
+        for kwarg in list(kwargs.keys()):
+            if kwarg not in do_bench_params:
+                del kwargs[kwarg]
+        try:
+            if "quantiles" in kwargs:
+                return self.triton_do_bench(_callable, **kwargs)[0]
+            elif "return_mode" in kwargs:
+                return self.triton_do_bench(_callable, **kwargs)
+            return self.triton_do_bench(_callable, **kwargs, return_mode="median")
+        except Exception as e:
+            # ErrorInvalidConfiguration
+            # Return inf to skip this config during autotuning
+            error_str = str(e).lower()
+            if "invalid configuration" in error_str:
+                logger.warning(
+                    "Skipping benchmark due to invalid configuration error: %s",
+                    error_str,
+                )
+                return float("inf")
+            raise
+
+
+class InductorBenchmarker(TritonBenchmarker):  # noqa: docstring_linter
+    @cached_property
+    def L2_cache_size(self: Self) -> int:
+        """Get the L2 cache size, in bytes, of the current device."""
+        device = torch.cuda.current_device()
+        props = torch.cuda.get_device_properties(device)
+        return props.L2_cache_size
+
+    def get_event_pairs(
+        self: Self, iters: int
+    ) -> list[tuple[torch.cuda.Event, torch.cuda.Event]]:
+        """Get `iters` pairs of CUDA events."""
+        return [
+            (
+                torch.cuda.Event(enable_timing=True),
+                torch.cuda.Event(enable_timing=True),
+            )
+            for _ in range(iters)
+        ]
+
+    def get_event_pairs_min_timing(
+        self: Self, event_pairs: list[tuple[torch.cuda.Event, torch.cuda.Event]]
+    ) -> float:
+        """Get the minimum timing, in milliseconds, for a group of CUDA event pairs."""
+        return min(
+            [
+                start_event.elapsed_time(end_event)
+                for start_event, end_event in event_pairs
+            ]
+        )
+
+    @may_distort_benchmarking_result
+    @time_and_count
+    def benchmark_gpu(  # type: ignore[override]
+        self: Self,
+        _callable: Callable[[], Any],
+        estimation_iters: int = 5,
+        memory_warmup_iters: int = 100,
+        benchmark_iters: int = 100,
+        max_benchmark_duration: int = 25,
+        return_mode: str = "min",
+        grad_to_none: list[torch.Tensor] | None = None,
+        is_vetted_benchmarking: bool = False,
+        **kwargs: Any,
+    ) -> float | list[float]:
+        """Benchmark a GPU callable using a custom benchmarking implementation.
+
+        Arguments:
+        - _callable: The callable to benchmark.
+
+        Keyword Arguments:
+        - estimation_iters: Optionally, the number of iterations to run `_callable`
+        during runtime estimation.
+        - memory_warmup_iters: Optionally, the number of iterations to flush the L2
+        cache before starting benchmarking.
+        - benchmark_iters: Optionally, the number of iterations to run `_callable`
+        during the benchmarking.
+        - max_benchmark_duration: Optionally, the maximum duration of the benchmarking,
+        in milliseconds. An estimated duration is calculated based on the values
+        of `memory_warmup_iters` and `benchmark_iters`, along with the estimated
+        runtime of `_callable` and various other factors, and we then shrink
+        `benchmark_iters` to fit in the allotted maximum duration.
+        - return_mode: Return mode for benchmark results. Options are "min" (default),
+        "all" (returns all measurements).
+        - grad_to_none: Optionally, a list of tensors whose gradients should be cleared
+        before each benchmark iteration.
+        - is_vetted_benchmarking: in deterministic mode, we only allow
+        benchmarking in vetted cases.
+        - **kwargs: Additional kwargs that may be passed to the fallback.
+
+        Returns:
+        - If return_mode="min": The minimum runtime of `_callable`, in milliseconds.
+        - If return_mode="all": List of all runtime measurements, in milliseconds.
+        """
+
+        if not is_vetted_benchmarking:
+            may_ban_benchmarking()
+
+        # we don't want any outside errors propagating into benchmarking
+        torch.cuda.synchronize()
+
+        # warmup `_callable` (and catches any failures in the process)
+        _callable()
+        torch.cuda.synchronize()
+
+        # see https://github.com/triton-lang/triton/pull/840 for why `dtype=torch.int`
+        buffer = torch.empty(self.L2_cache_size // 4, dtype=torch.int, device="cuda")
+        buffer.zero_()
+
+        # estimate the runtime of `_callable`
+        event_pairs = self.get_event_pairs(estimation_iters)
+        for start_event, end_event in event_pairs:
+            # Clear gradients before timing (matches triton.testing.do_bench)
+            if grad_to_none is not None:
+                for x in grad_to_none:
+                    x.grad = None
+            buffer.zero_()
+            start_event.record()
+            _callable()
+            end_event.record()
+        torch.cuda.synchronize()
+        estimated_timing = self.get_event_pairs_min_timing(event_pairs)
+
+        # adjust `benchmark_iters` to fit in the maximum benchmarking duration
+        if estimated_timing > 0:
+            benchmark_iters = max(
+                min(benchmark_iters, int(max_benchmark_duration // estimated_timing)), 1
+            )
+
+        # do the memory warmup
+        for _ in range(memory_warmup_iters):
+            buffer.zero_()
+
+        # benchmark `_callable`
+        event_pairs = self.get_event_pairs(benchmark_iters)
+        for start_event, end_event in event_pairs:
+            # Clear gradients before timing (matches triton.testing.do_bench)
+            if grad_to_none is not None:
+                for x in grad_to_none:
+                    x.grad = None
+            buffer.zero_()
+            start_event.record()
+            _callable()
+            end_event.record()
+        torch.cuda.synchronize()
+
+        # explicitly delete the buffer, sometimes helps memory
+        # footprint metrics in OSS Inductor performance benchmarks
+        del buffer
+
+        # Return based on the requested mode
+        if return_mode == "all":
+            # Get all timings from event pairs
+            all_timings = [
+                start_event.elapsed_time(end_event)
+                for start_event, end_event in event_pairs
+            ]
+            return all_timings
+        elif return_mode == "min":
+            benchmarked_timing = self.get_event_pairs_min_timing(event_pairs)
+            # return the minimum of `estimated_timing` and `benchmarked_timing`,
+            # we just want the minimum timing overall so we might as well check both
+            return min(estimated_timing, benchmarked_timing)
+        else:
+            raise ValueError(
+                f"Unsupported return_mode: {return_mode}. Use 'min' or 'all'."
+            )
+
+
+benchmarker = (
+    InductorBenchmarker() if use_experimental_benchmarker else TritonBenchmarker()
+)

outputs/audit_venv/lib/python3.11/site-packages/torch/_inductor/runtime/cache_dir_utils.py

@@ -0,0 +1,54 @@
+import getpass
+import os
+import re
+import tempfile
+from collections.abc import Generator
+from contextlib import contextmanager
+
+from torch._environment import is_fbcode
+
+
+# Factoring out to file without torch dependencies
+
+
+def cache_dir() -> str:
+    cache_dir = os.environ.get("TORCHINDUCTOR_CACHE_DIR")
+    if cache_dir is None:
+        os.environ["TORCHINDUCTOR_CACHE_DIR"] = cache_dir = default_cache_dir()
+    os.makedirs(cache_dir, exist_ok=True)
+    return cache_dir
+
+
+def default_cache_dir() -> str:
+    sanitized_username = re.sub(r'[\\/:*?"<>|]', "_", getpass.getuser())
+    return os.path.join(
+        tempfile.gettempdir() if not is_fbcode() else "/var/tmp",
+        "torchinductor_" + sanitized_username,
+    )
+
+
+def triton_cache_dir(device: int) -> str:
+    if (directory := os.getenv("TRITON_CACHE_DIR")) is not None:
+        return directory
+    return os.path.join(
+        cache_dir(),
+        "triton",
+        str(device),
+    )
+
+
+@contextmanager
+def temporary_cache_dir(directory: str) -> Generator[None, None, None]:
+    from torch._inductor.utils import clear_caches
+
+    original = os.environ.get("TORCHINDUCTOR_CACHE_DIR")
+    os.environ["TORCHINDUCTOR_CACHE_DIR"] = directory
+    try:
+        clear_caches()
+        yield
+    finally:
+        clear_caches()
+        if original is None:
+            del os.environ["TORCHINDUCTOR_CACHE_DIR"]
+        else:
+            os.environ["TORCHINDUCTOR_CACHE_DIR"] = original

outputs/audit_venv/lib/python3.11/site-packages/torch/_inductor/runtime/caching/__init__.py

@@ -0,0 +1,32 @@
+from . import config, encoders, memoizers
+from .context import IsolationSchema, SelectedCompileContext, SelectedRuntimeContext
+from .exceptions import (
+    CacheError,
+    FileLockTimeoutError,
+    KeyEncodingError,
+    LockTimeoutError,
+    SystemError,
+    UserError,
+    ValueDecodingError,
+    ValueEncodingError,
+)
+from .interfaces import Memoizer, PersistentMemoizer
+
+
+__all__ = [
+    "CacheError",
+    "FileLockTimeoutError",
+    "IsolationSchema",
+    "KeyEncodingError",
+    "LockTimeoutError",
+    "Memoizer",
+    "PersistentMemoizer",
+    "SelectedCompileContext",
+    "SelectedRuntimeContext",
+    "SystemError",
+    "UserError",
+    "ValueDecodingError",
+    "ValueEncodingError",
+    "encoders",
+    "memoizers",
+]

outputs/audit_venv/lib/python3.11/site-packages/torch/_inductor/runtime/caching/config.py

@@ -0,0 +1,167 @@
+import os
+from collections.abc import Callable
+from functools import cache, partial
+from typing import TypeVar
+
+import torch
+from torch._environment import is_fbcode
+
+
+T = TypeVar("T")
+
+
+def _is_force_disable_caches() -> bool:
+    """Check if caching is force disabled via inductor config.
+
+    This defers importing torch._inductor.config to avoid circular imports.
+
+    Returns:
+        True if force_disable_caches is set in inductor config, False otherwise.
+    """
+    from torch._inductor import config as inductor_config
+
+    return inductor_config.force_disable_caches
+
+
+@cache
+def _env_var_val(env_var: str, default: T) -> str | T:
+    """Get the value of an environment variable or return the default.
+
+    Args:
+        env_var: Environment variable name to check
+        default: Default value to return if environment variable is not set
+
+    Returns:
+        The value from the environment variable as a string, or the default if not set
+    """
+    return os.environ.get(env_var, default)
+
+
+@cache
+def _env_var_config(env_var: str, default: bool) -> bool:
+    env_val = _env_var_val(env_var, None)
+    if env_val is not None:
+        return env_val == "1"
+    return default
+
+
+@cache
+def _versioned_config(
+    jk_name: str,
+    this_version: int,
+    oss_default: bool,
+    env_var_override: str | None = None,
+) -> bool:
+    """
+    A versioned configuration utility that determines boolean settings based on:
+    1. Environment variable override (highest priority)
+    2. JustKnobs version comparison in fbcode environments
+    3. OSS default fallback
+
+    This function enables gradual rollouts of features in fbcode by comparing
+    a local version against a JustKnobs-controlled remote version, while
+    allowing environment variable overrides for testing and OSS defaults
+    for non-fbcode environments.
+
+    Args:
+        jk_name: JustKnobs key name (e.g., "pytorch/inductor:feature_version")
+        this_version: Local version number to compare against JustKnobs version
+        oss_default: Default value to use in non-fbcode environments
+        env_var_override: Optional environment variable name that, when set,
+                         overrides all other logic
+
+    Returns:
+        bool: Configuration value determined by the priority order above
+    """
+    if (
+        env_var_override
+        and (env_var_value := os.environ.get(env_var_override)) is not None
+    ):
+        return env_var_value == "1"
+    elif is_fbcode():
+        # this method returns 0 on failure, which we should check for specifically.
+        # in the case of JK failure, the safe bet is to simply disable the config
+        jk_version: int = torch._utils_internal.justknobs_getval_int(jk_name)
+        return (this_version >= jk_version) and (jk_version != 0)
+    return oss_default
+
+
+# toggles the entire caching module, but only when calling through the
+# public facing interfaces. get/insert operations become no-ops in the sense
+# that get will always miss and insert will never insert; record becomes a
+# no-op in the sense that the function will always be called and the cache
+# will never be accessed
+_CACHING_MODULE_VERSION: int = 0
+_CACHING_MODULE_VERSION_JK: str = "pytorch/inductor:caching_module_version"
+_CACHING_MODULE_OSS_DEFAULT: bool = False
+_CACHING_MODULE_ENV_VAR_OVERRIDE: str = "TORCHINDUCTOR_ENABLE_CACHING_MODULE"
+
+
+def _is_caching_module_enabled_base() -> bool:
+    """Base check for caching module enablement via versioned config."""
+    return _versioned_config(
+        _CACHING_MODULE_VERSION_JK,
+        _CACHING_MODULE_VERSION,
+        _CACHING_MODULE_OSS_DEFAULT,
+        _CACHING_MODULE_ENV_VAR_OVERRIDE,
+    )
+
+
+def IS_CACHING_MODULE_ENABLED() -> bool:
+    """Check if the caching module is enabled.
+
+    Returns False if:
+    - The versioned config disables it
+    - force_disable_caches is set in inductor config
+
+    Returns:
+        True if caching module is enabled, False otherwise.
+    """
+    if not _is_caching_module_enabled_base():
+        return False
+    if _is_force_disable_caches():
+        return False
+    return True
+
+
+# Controls whether the Memoizer dumps its cache to a JSON file on destruction.
+# When enabled, the Memoizer will write its in-memory cache to a JSON file
+# (memoizer_cache.json in the cache directory) on destruction. This dump file
+# can later be used to pre-populate Memoizers via CACHE_DUMP_FILE_PATH.
+#
+# This is useful for:
+# - Debugging and inspection of cached values
+# - Creating cache snapshots that can be reused across runs
+# - Pre-warming caches for subsequent executions
+#
+# Set via environment variable: TORCHINDUCTOR_DUMP_MEMOIZER_CACHE=1
+_DUMP_MEMOIZER_CACHE_ENV_VAR: str = "TORCHINDUCTOR_DUMP_MEMOIZER_CACHE"
+_DUMP_MEMOIZER_CACHE_DEFAULT: bool = False
+IS_DUMP_MEMOIZER_CACHE_ENABLED: Callable[[], bool] = partial(
+    _env_var_config,
+    _DUMP_MEMOIZER_CACHE_ENV_VAR,
+    _DUMP_MEMOIZER_CACHE_DEFAULT,
+)
+
+
+# Path to a cache dump file to pre-populate Memoizers on initialization.
+# This should point to a JSON file produced by IS_DUMP_MEMOIZER_CACHE_ENABLED.
+#
+# When set, Memoizers will load cached entries from this file on initialization,
+# allowing cache values to be reused across separate runs without recomputation.
+# This is particularly useful for:
+# - Pre-warming caches with known-good values
+# - Reproducing behavior from a previous run
+# - Sharing cached computations across different environments
+#
+# The dump file format is produced by the Memoizer's _dump_to_disk method when
+# IS_DUMP_MEMOIZER_CACHE_ENABLED is set to true.
+#
+# Set via environment variable: TORCHINDUCTOR_CACHE_DUMP_FILE_PATH=/path/to/dump.json
+_CACHE_DUMP_FILE_PATH_ENV_VAR: str = "TORCHINDUCTOR_CACHE_DUMP_FILE_PATH"
+_CACHE_DUMP_FILE_PATH_DEFAULT: str | None = None
+CACHE_DUMP_FILE_PATH: Callable[[], str | None] = partial(
+    _env_var_val,
+    _CACHE_DUMP_FILE_PATH_ENV_VAR,
+    _CACHE_DUMP_FILE_PATH_DEFAULT,
+)

outputs/audit_venv/lib/python3.11/site-packages/torch/_inductor/runtime/caching/context.py

@@ -0,0 +1,310 @@
+"""Context management for PyTorch Inductor runtime caching.
+
+This module provides context classes for collecting configuration and environment
+information used in caching decisions for PyTorch's Inductor runtime.
+"""
+
+import json
+from abc import ABC, abstractmethod
+from base64 import b64encode
+from collections.abc import Sequence
+from functools import cache
+from hashlib import sha256
+from typing_extensions import override, TypedDict
+
+import torch
+
+
+class _Context(ABC):
+    """Abstract base class for context providers.
+
+    Context providers collect specific configuration and environment information
+    that affects compilation and runtime behavior.
+    """
+
+    @staticmethod
+    @abstractmethod
+    def forms_of_context() -> Sequence[str]:
+        """Return a sequence of context form names provided by this context class.
+
+        Returns:
+            A sequence of strings representing the available context forms.
+        """
+
+
+class _RuntimeContext(_Context):
+    """Context provider for runtime configuration and environment settings.
+
+    Collects configuration settings that affect runtime behavior but not
+    compilation, such as Inductor configs, determinism settings, and CUDA
+    matmul precision configurations.
+    """
+
+    @override
+    @staticmethod
+    def forms_of_context() -> Sequence[str]:
+        """Return the runtime context forms provided by this class.
+
+        Returns:
+            A sequence containing the available runtime context forms:
+            - "inductor_configs": PyTorch Inductor configuration settings
+            - "torch_determinism_configs": Deterministic algorithm settings
+            - "cuda_matmul_precision_configs": CUDA matrix multiplication precision settings
+        """
+        return (
+            "inductor_configs",
+            "torch_determinism_configs",
+            "cuda_matmul_precision_configs",
+        )
+
+    @staticmethod
+    def inductor_configs() -> dict[str, object]:
+        """Get portable Inductor configuration settings.
+
+        Returns:
+            A dictionary containing Inductor configuration settings,
+            including private configs.
+        """
+        from torch._inductor import config
+
+        return config.save_config_portable(ignore_private_configs=False)
+
+    @staticmethod
+    def torch_determinism_configs() -> dict[str, object]:
+        """Get PyTorch deterministic algorithm configuration settings.
+
+        Returns:
+            A dictionary containing deterministic algorithm settings:
+            - Whether deterministic algorithms are enabled
+            - Whether deterministic algorithm warnings are enabled
+            - Fill uninitialized memory setting
+        """
+        return {
+            "torch.are_deterministic_algorithms_enabled": torch.are_deterministic_algorithms_enabled(),
+            "torch.is_deterministic_algorithms_warn_only_enabled": (
+                torch.is_deterministic_algorithms_warn_only_enabled()
+            ),
+            "torch.utils.deterministic.fill_uninitialized_memory": getattr(
+                torch.utils.deterministic, "fill_uninitialized_memory", None
+            ),
+        }
+
+    @staticmethod
+    def cuda_matmul_precision_configs() -> dict[str, object]:
+        """Get CUDA matrix multiplication precision configuration settings.
+
+        Returns:
+            A dictionary containing CUDA matmul precision settings:
+            - FP32 precision setting
+            - FP16 reduced precision reduction allowance
+            - BF16 reduced precision reduction allowance
+        """
+        return {
+            "torch.backends.cuda.matmul.fp32_precision": torch.backends.cuda.matmul.fp32_precision,
+            "torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction": (
+                torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction
+            ),
+            "torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction": (
+                torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction
+            ),
+        }
+
+
+class _CompileContext(_Context):
+    """Context provider for compilation-related configuration and environment settings.
+
+    Collects information that affects compilation behavior, such as PyTorch and Triton
+    versions, runtime environment, and accelerator properties.
+    """
+
+    @override
+    @staticmethod
+    def forms_of_context() -> Sequence[str]:
+        """Return the compile context forms provided by this class.
+
+        Returns:
+            A sequence containing the available compile context forms:
+            - "torch_version_hash": PyTorch version hash
+            - "triton_version_hash": Triton version hash (if available)
+            - "runtime": Runtime type (CUDA/HIP/None)
+            - "runtime_version": Runtime version string
+            - "accelerator_properties": GPU/accelerator properties
+        """
+        return (
+            "torch_version_hash",
+            "triton_version_hash",
+            "runtime",
+            "runtime_version",
+            "accelerator_properties",
+        )
+
+    @cache
+    @staticmethod
+    def torch_version_hash() -> str:
+        """Get base64-encoded PyTorch version hash.
+
+        Returns:
+            A base64-encoded string representing the PyTorch version hash.
+        """
+        from torch._inductor.codecache import torch_key
+
+        return b64encode(torch_key()).decode()
+
+    @cache
+    @staticmethod
+    def triton_version_hash() -> str | None:
+        """Get Triton version key if Triton is available.
+
+        Returns:
+            Triton version key if Triton is available, None otherwise.
+        """
+        from torch._inductor.runtime.triton_compat import HAS_TRITON, triton_key
+
+        return triton_key() if HAS_TRITON else None
+
+    @cache
+    @staticmethod
+    def runtime() -> str | None:
+        """Determine the runtime type based on available backends.
+
+        Returns:
+            "CUDA" if CUDA is available, "HIP" if HIP is available, None otherwise.
+        """
+        return "CUDA" if torch.version.cuda else "HIP" if torch.version.hip else None
+
+    @cache
+    @staticmethod
+    def runtime_version() -> str | None:
+        """Get the version string for the detected runtime.
+
+        Returns:
+            Version string for the current runtime (CUDA or HIP), or None if
+            no supported runtime is detected.
+        """
+        return {
+            "CUDA": torch.version.cuda,
+            "HIP": torch.version.hip,
+            "None": None,
+        }.get(_CompileContext.runtime() or "None")
+
+    @cache
+    @staticmethod
+    def accelerator_properties() -> str | None:
+        """Get string representation of CUDA device properties.
+
+        Returns:
+            String representation of CUDA device properties if a runtime is
+            available, None otherwise.
+        """
+        return (
+            repr(torch.cuda.get_device_properties())
+            if _CompileContext.runtime() and torch.cuda.is_available()
+            else None
+        )
+
+
+class SelectedRuntimeContext(TypedDict):
+    inductor_configs: bool
+    torch_determinism_configs: bool
+    cuda_matmul_precision_configs: bool
+
+
+class SelectedCompileContext(TypedDict):
+    torch_version_hash: bool
+    triton_version_hash: bool
+    runtime: bool
+    runtime_version: bool
+    accelerator_properties: bool
+
+
+class IsolationSchema(TypedDict):
+    """Schema for specifying which context forms to include in cache isolation.
+
+    Attributes:
+        runtime_context: Either True (include all runtime context), False (exclude all),
+                        or a SelectedRuntimeContext dict specifying which forms to include.
+        compile_context: Either True (include all compile context), False (exclude all),
+                        or a SelectedCompileContext dict specifying which forms to include.
+    """
+
+    runtime_context: SelectedRuntimeContext | bool
+    compile_context: SelectedCompileContext | bool
+
+
+_DEFAULT_ISOLATION_SCHEMA: IsolationSchema = IsolationSchema(
+    runtime_context=True, compile_context=True
+)
+
+
+def _collect_runtime_context(
+    selection: SelectedRuntimeContext | bool,
+) -> dict[str, object] | None:
+    """Collect runtime context based on selection.
+
+    Args:
+        selection: True to include all, False to exclude all, or a dict
+                  specifying which forms to include.
+
+    Returns:
+        Dictionary of selected context data, or None if excluded.
+    """
+    return {
+        form: getattr(_RuntimeContext, form)()
+        for form in _RuntimeContext.forms_of_context()
+        if selection is True or (selection and selection.get(form, False))
+    } or None
+
+
+def _collect_compile_context(
+    selection: SelectedCompileContext | bool,
+) -> dict[str, object] | None:
+    """Collect compile context based on selection.
+
+    Args:
+        selection: True to include all, False to exclude all, or a dict
+                  specifying which forms to include.
+
+    Returns:
+        Dictionary of selected context data, or None if excluded.
+    """
+    return {
+        form: getattr(_CompileContext, form)()
+        for form in _CompileContext.forms_of_context()
+        if selection is True or (selection and selection.get(form, False))
+    } or None
+
+
+def _isolation_context(
+    ischema: IsolationSchema = _DEFAULT_ISOLATION_SCHEMA,
+) -> dict[str, object]:
+    """Generate context data based on the isolation schema.
+
+    Args:
+        ischema: Schema specifying which context forms to include.
+                Defaults to including all runtime and compile context.
+
+    Returns:
+        A dictionary containing the selected context data with keys
+        "runtime_context" and "compile_context", where each value is
+        either None (if excluded) or a dict of context form data.
+    """
+    return {
+        "runtime_context": _collect_runtime_context(ischema["runtime_context"]),
+        "compile_context": _collect_compile_context(ischema["compile_context"]),
+    }
+
+
+def _isolation_key(ischema: IsolationSchema = _DEFAULT_ISOLATION_SCHEMA) -> str:
+    """Generate a unique key for the given isolation schema.
+
+    Args:
+        ischema: Schema specifying which context forms to include.
+                Defaults to including all runtime and compile context.
+
+    Returns:
+        A 32-character hexadecimal string that uniquely identifies
+        the context specified by the isolation schema.
+    """
+    return sha256(
+        json.dumps(_isolation_context(ischema), sort_keys=True).encode()
+    ).hexdigest()[:32]

outputs/audit_venv/lib/python3.11/site-packages/torch/_inductor/runtime/caching/encoders.py

@@ -0,0 +1,68 @@
+"""
+Custom encoder functions
+
+This module provides reusable encoder functions that convert function parameters
+into JSON-serializable dictionaries for caching purposes.
+"""
+
+from typing_extensions import TypedDict
+
+import torch
+from torch import Tensor
+from torch._inductor.pattern_matcher import Match
+from torch._inductor.runtime.caching.utils import _encode_tensor, EncodedTensor
+
+
+class ShouldPadEncodedParams(TypedDict):
+    """TypedDict for encoded should_pad parameters."""
+
+    mat1: EncodedTensor
+    mat2: EncodedTensor
+    op: str
+    input: EncodedTensor | None
+    mat1_exclude_padding_time: bool
+    mat2_exclude_padding_time: bool
+    tf32: bool
+
+
+def should_pad_params_encoder(
+    match: Match,
+    mat1: Tensor,
+    mat2: Tensor,
+    op: torch._ops.OpOverloadPacket,
+    input: Tensor | None = None,
+) -> ShouldPadEncodedParams:
+    """Encode parameters for _should_pad into a human-readable dict.
+
+    This encoder extracts only the information needed for caching:
+    - Tensor shape, stride, and dtype (not the actual data)
+    - Whether padding time should be excluded for mat1 and mat2
+    - The operation as a string
+
+    Args:
+        match: The pattern match object
+        mat1: First matrix tensor
+        mat2: Second matrix tensor
+        op: The operation being performed
+        input: Optional input tensor for addmm
+
+    Returns:
+        A dict containing the encoded parameters in human-readable form
+    """
+    # Import here to avoid circular dependency
+    from torch._inductor.fx_passes.pad_mm import should_exclude_padding_time
+
+    return ShouldPadEncodedParams(
+        mat1=_encode_tensor(mat1),
+        mat2=_encode_tensor(mat2),
+        op=str(op),
+        input=_encode_tensor(input) if input is not None else None,
+        mat1_exclude_padding_time=should_exclude_padding_time(match, "mat1"),
+        mat2_exclude_padding_time=should_exclude_padding_time(match, "mat2"),
+        tf32=False
+        if mat1.dtype != torch.float32
+        else bool(
+            torch.backends.cuda.matmul.fp32_precision == "tf32"
+            or torch.backends.mkldnn.fp32_precision == "tf32"
+        ),
+    )

outputs/audit_venv/lib/python3.11/site-packages/torch/_inductor/runtime/caching/exceptions.py

@@ -0,0 +1,96 @@
+"""Exception classes for PyTorch Inductor runtime caching.
+
+This module defines a hierarchy of exceptions used throughout the caching system.
+All custom exceptions inherit from CacheError, with UserError serving as a base
+for user-facing errors that also inherit from TypeError for compatibility.
+"""
+
+from threading import Lock
+
+from filelock import BaseFileLock
+
+
+class CacheError(Exception):
+    """Base class for all caching-related errors.
+
+    This is the root exception class for all custom exceptions raised by the caching
+    module, providing a common interface for error handling and logging.
+    """
+
+
+class SystemError(CacheError, RuntimeError):
+    """Base class for system-level caching errors.
+
+    This class represents errors that occur during cache operations, such as
+    storage or retrieval failures. It inherits from RuntimeError to indicate
+    that the error is not caused by user input.
+    """
+
+
+class LockTimeoutError(SystemError):
+    """Error raised when a lock operation times out.
+
+    This exception is raised when a lock operation exceeds the specified timeout
+    limit, indicating that the lock could not be acquired within the allotted time.
+    """
+
+    def __init__(self, lock: Lock, timeout: float) -> None:
+        """Initialize the lock timeout error with detailed lock information.
+
+        Args:
+            lock: The lock object that timed out.
+            timeout: The timeout limit that was exceeded.
+        """
+        super().__init__(f"Failed to acquire lock {lock} within {timeout} seconds.")
+
+
+class FileLockTimeoutError(SystemError):
+    """Error raised when a file lock operation times out.
+
+    This exception is raised when a file lock operation exceeds the specified timeout
+    limit, indicating that the lock could not be acquired within the allotted time.
+    """
+
+    def __init__(self, flock: BaseFileLock, timeout: float) -> None:
+        """Initialize the file lock timeout error with detailed lock information.
+
+        Args:
+            flock: The file lock object that timed out.
+            timeout: The timeout limit that was exceeded.
+        """
+        super().__init__(
+            f"Failed to acquire file lock {flock} within {timeout} seconds."
+        )
+
+
+class UserError(CacheError, TypeError):
+    """Base class for user-facing cache errors that also inherit from TypeError.
+
+    This class combines CacheError with TypeError to provide compatibility
+    with existing exception handling patterns while maintaining the cache
+    error hierarchy. All user-facing cache errors should inherit from this class.
+    """
+
+
+class KeyEncodingError(UserError):
+    """Base class for errors that occur during cache key encoding operations.
+
+    Raised when cache keys cannot be properly encoded for storage or transmission.
+    This includes serialization, hashing, or other encoding-related failures.
+    """
+
+
+class ValueEncodingError(UserError):
+    """Base class for errors that occur during cache value encoding operations.
+
+    Raised when cache values cannot be properly encoded for storage or transmission.
+    This includes serialization, compression, or other encoding-related failures.
+    """
+
+
+class ValueDecodingError(UserError):
+    """Base class for errors that occur during cache value decoding operations.
+
+    Raised when cached values cannot be properly decoded during retrieval.
+    This includes deserialization, decompression, or other decoding-related failures.
+    """

outputs/audit_venv/lib/python3.11/site-packages/torch/_inductor/runtime/caching/implementations.py

@@ -0,0 +1,422 @@
+"""Cache implementation classes
+
+This module provides concrete implementations of caching backends including
+in-memory, on-disk, and remote caching strategies. Each implementation follows
+the abstract _CacheImpl interface and provides thread-safe operations with
+appropriate locking mechanisms.
+"""
+
+from abc import ABC, abstractmethod
+from collections.abc import Generator
+from contextlib import contextmanager
+from dataclasses import dataclass
+from hashlib import sha256
+from io import BufferedReader, BufferedWriter
+from os import PathLike
+from pathlib import Path
+from threading import Lock
+from typing import Generic, TypeVar
+from typing_extensions import override
+
+from filelock import BaseFileLock, FileLock
+
+from . import locks
+
+
+# Type variable for cache value types
+_V = TypeVar("_V")
+
+
+@dataclass
+class Hit(Generic[_V]):
+    """Result wrapper for hits on cache get operations.
+
+    Allows distinguishing between a cache miss and a cached None value.
+
+    Attributes:
+        value: The cached value. The type depends on the cache implementation:
+              - InMemoryCache: object (any Python object)
+              - OnDiskCache: bytes
+              - RemoteCache: bytes
+    """
+
+    value: _V
+
+
+class Miss:
+    """Sentinel class representing a cache miss.
+
+    Used to distinguish between a cached None value and a cache miss
+    when None is a valid cached value.
+    """
+
+
+# Singleton instance for cache miss sentinel
+miss = Miss()
+
+
+class _CacheImpl(ABC, Generic[_V]):
+    """Abstract base class for cache implementations.
+
+    This class defines the interface that all cache implementations must follow.
+    It provides thread-safe operations through a locking mechanism and supports
+    both get and insert operations.
+
+    Type Parameters:
+        _V: The type of values stored in the cache. Each implementation specifies
+           its own value type (e.g., object for in-memory, bytes for on-disk).
+    """
+
+    def __init__(self) -> None:
+        """Initialize the cache implementation with a threading lock."""
+        self._lock: Lock = Lock()
+
+    @property
+    def lock(self) -> locks._LockProtocol:
+        """Get a context manager for acquiring the cache lock.
+
+        Locking of the cache is not done by the implementation itself, but by the
+        interface that uses it. The interface may want to hold the lock for longer
+        than a single cache operation, for example when dealing with multiple
+        cache implementations at once, so we leave that decision up to the interface.
+
+        Args:
+            timeout: Optional timeout in seconds (float) for acquiring the lock.
+
+        Returns:
+            A callable that returns a context manager for the lock.
+        """
+
+        def _lock_with_timeout(
+            timeout: float | None = None,
+        ) -> locks._LockContextManager:
+            return locks._acquire_lock_with_timeout(self._lock, timeout)
+
+        return _lock_with_timeout
+
+    @abstractmethod
+    def get(self, key: str) -> Hit[_V] | None:
+        """Retrieve a value from the cache.
+
+        Args:
+            key: The key to look up in the cache (must be str).
+
+        Returns:
+            A Hit object on cache hit where Hit.value is the cached value,
+            or None on cache miss.
+        """
+
+    @abstractmethod
+    def insert(self, key: str, value: _V) -> bool:
+        """Insert a key-value pair into the cache.
+
+        Args:
+            key: The key to insert (must be str).
+            value: The value to associate with the key.
+
+        Returns:
+            True if the insertion was successful, False if not inserted.
+        """
+
+
+class _InMemoryCacheImpl(_CacheImpl[_V], Generic[_V]):
+    """In-memory cache implementation using a dictionary.
+
+    This implementation stores key-value pairs in a Python dictionary,
+    with keys being pickled for consistent hashing. It provides fast
+    access but is limited by available memory and process lifetime.
+
+    Type Parameters:
+        _V: The type of values stored in the cache.
+    """
+
+    def __init__(self) -> None:
+        """Initialize the in-memory cache with an empty dictionary."""
+        super().__init__()
+        self._memory: dict[str, _V] = {}
+
+    @override
+    def get(self, key: str) -> Hit[_V] | None:
+        """Retrieve a value from the in-memory cache.
+
+        Args:
+            key: The key to look up (must be str).
+
+        Returns:
+            A Hit object on cache hit where Hit.value is the cached value,
+            or None on cache miss.
+        """
+        if key in self._memory:
+            return Hit(value=self._memory[key])
+        return None
+
+    @override
+    def insert(self, key: str, value: _V) -> bool:
+        """Insert a key-value pair into the in-memory cache.
+
+        Args:
+            key: The key to insert (must be str).
+            value: The value to associate with the key.
+
+        Returns:
+            True if the insertion was successful (key was new),
+            False if not inserted (key already existed).
+        """
+        if key not in self._memory:
+            self._memory[key] = value
+            return True
+        return False
+
+
+class _OnDiskCacheImpl(_CacheImpl[bytes]):
+    """On-disk cache implementation using file system storage.
+
+    This implementation stores cached data as files on disk, with version
+    headers to handle cache invalidation. It uses file locking to ensure
+    thread safety across processes and provides persistent storage that
+    survives process restarts.
+
+    The value type is 'bytes' since data must be serialized to bytes for disk storage.
+
+    Attributes:
+        _version: _Version number for cache format compatibility.
+        _version_header_length: Length of the version header in bytes.
+    """
+
+    _version: int = 0
+    _version_header_length: int = 4
+
+    def __init__(self, sub_dir: PathLike[str] | None = None) -> None:
+        """Initialize the on-disk cache with a specified subdirectory.
+
+        Args:
+            sub_dir: Subdirectory name within the cache directory.
+                    Defaults to empty string if not specified.
+        """
+        self._cache_dir: Path = self._base_dir / (sub_dir or "")
+        self._flock: BaseFileLock = FileLock(str(self._cache_dir / "dir.lock"))
+
+    @property
+    def _base_dir(self) -> Path:
+        """Get the base directory for cache storage.
+
+        Returns:
+            Path to the cache directory based on the default cache dir
+            and the specified subdirectory.
+        """
+        from torch._inductor.runtime.runtime_utils import default_cache_dir
+
+        return Path(default_cache_dir(), "cache")
+
+    def _fpath_from_key(self, key: str) -> Path:
+        """Generate a file path from a cache key.
+
+        Args:
+            key: The cache key to convert to a file path (must be str).
+
+        Returns:
+            A Path object representing the file location for this key.
+        """
+        return self._cache_dir / key
+
+    @classmethod
+    def _version_header(cls) -> bytes:
+        """Generate the version header bytes.
+
+        Returns:
+            A byte string representing the current cache version header.
+        """
+        return sha256(str(cls._version).encode()).digest()[: cls._version_header_length]
+
+    def _version_header_matches(self, fp: BufferedReader) -> bool:
+        """Check if the file's version header matches the current version.
+
+        Args:
+            fp: File pointer positioned at the start of the file.
+
+        Returns:
+            True if the version header matches, False otherwise.
+        """
+        return fp.read(self._version_header_length) == self._version_header()
+
+    def _write_version_header(self, fp: BufferedWriter) -> None:
+        """Write the version header to a file.
+
+        Args:
+            fp: File pointer where the version header should be written.
+        """
+        fp.write(self._version_header())
+
+    @override
+    @property
+    def lock(self) -> locks._LockProtocol:
+        """Get a context manager for acquiring the file lock.
+
+        Uses file locking to ensure thread safety across processes.
+
+        Args:
+            timeout: Optional timeout in seconds (float) for acquiring the file lock.
+
+        Returns:
+            A callable that returns a context manager for the file lock.
+        """
+
+        def _lock_with_timeout(
+            timeout: float | None = None,
+        ) -> locks._LockContextManager:
+            return locks._acquire_flock_with_timeout(self._flock, timeout)
+
+        return _lock_with_timeout
+
+    @override
+    def get(self, key: str) -> Hit[bytes] | None:
+        """Retrieve a value from the on-disk cache.
+
+        Args:
+            key: The key to look up in the cache (must be str).
+
+        Returns:
+            A Hit object on cache hit where Hit.value is the cached value (bytes),
+            or None on cache miss or version mismatch.
+        """
+        fpath: Path = self._fpath_from_key(key)
+
+        if not fpath.is_file():
+            return None
+
+        pickled_value: bytes | None = None
+        with open(fpath, "rb") as fp:
+            if self._version_header_matches(fp):
+                pickled_value = fp.read()
+
+        if not pickled_value:
+            # if pickled_value is still None, even though the file exists, then
+            # we know that the version header did not match. in this case implementation
+            # is up to preference, we choose to remove entries that do not match
+            # the version header so that the key can be re-cached later with the correct
+            # version header
+            fpath.unlink()
+            return None
+
+        return Hit(value=pickled_value)
+
+    @override
+    def insert(self, key: str, value: bytes) -> bool:
+        """Insert a key-value pair into the on-disk cache.
+
+        Args:
+            key: The key to insert (must be str).
+            value: The value to associate with the key (must be bytes).
+
+        Returns:
+            True if successfully inserted, False if the key already exists
+            with a valid version.
+        """
+        fpath: Path = self._fpath_from_key(key)
+        fpath.parent.mkdir(parents=True, exist_ok=True)
+
+        r_fp, w_fp, inserted = None, None, False
+        try:
+            w_fp = open(fpath, "xb")  # noqa: SIM115
+        except FileExistsError:
+            is_stale: bool = False
+            with open(fpath, "rb") as r_fp:
+                is_stale = not self._version_header_matches(r_fp)
+
+            if is_stale:
+                # same story as above, in this case the version header doesn't
+                # match so we choose to remove the old entry so that the new
+                # k/v pair can be cached
+                fpath.unlink()
+                w_fp = open(fpath, "xb")  # noqa: SIM115
+            else:
+                w_fp = None
+        finally:
+            if w_fp:
+                try:
+                    self._write_version_header(w_fp)
+                    w_fp.write(value)
+                    inserted = True
+                finally:
+                    w_fp.close()
+
+        return inserted
+
+
+try:
+    from .fb.implementations import _RemoteCacheImpl
+except ModuleNotFoundError:
+
+    class _RemoteCacheImpl(_CacheImpl[bytes]):
+        """Fallback remote cache implementation for non-Facebook environments.
+
+        This is a no-op implementation that always raises NotImplementedError.
+        The actual remote cache implementation is provided in the `.fb` module
+        for Facebook-specific environments.
+
+        The value type is 'bytes' for consistency with the FB implementation.
+
+        Attributes:
+            _version: _Version number for cache format compatibility.
+            has_strong_consistency: Whether the remote cache provides strong
+                                   consistency guarantees.
+        """
+
+        _version: int = 0
+        has_strong_consistency: bool = False
+
+        def __init__(self) -> None:
+            """Initialize the fallback remote cache implementation.
+
+            Note: We don't need to initialize any form of lock since this
+            implementation provides a pseudo-lock context manager.
+            """
+
+        @override
+        @property
+        def lock(self) -> locks._LockProtocol:
+            """Get a pseudo lock that does nothing.
+
+            Most remote cache implementations don't have an ability to implement
+            any form of locking, so we provide a no-op pseudo-lock for consistency
+            with the interface.
+
+            Args:
+                timeout: Optional timeout in seconds (float). Ignored in this
+
+            Returns:
+                A callable that returns a no-op context manager.
+            """
+
+            @contextmanager
+            def pseudo_lock(
+                timeout: float | None = None,
+            ) -> Generator[None, None, None]:
+                yield
+
+            return pseudo_lock
+
+        @override
+        def get(self, key: str) -> Hit[bytes] | None:
+            """Raise NotImplementedError for remote cache get operations.
+
+            Args:
+                key: The key to look up (must be str, ignored).
+
+            Raises:
+                NotImplementedError: Always raised as this is a fallback implementation.
+            """
+            raise NotImplementedError
+
+        @override
+        def insert(self, key: str, value: bytes) -> bool:
+            """Raise NotImplementedError for remote cache insert operations.
+
+            Args:
+                key: The key to insert (must be str, ignored).
+                value: The value to insert (must be bytes, ignored).
+
+            Raises:
+                NotImplementedError: Always raised as this is a fallback implementation.
+            """
+            raise NotImplementedError

outputs/audit_venv/lib/python3.11/site-packages/torch/_inductor/runtime/caching/interfaces.py

@@ -0,0 +1,1280 @@
+"""Public interfaces for PyTorch Inductor runtime caching.
+
+This module provides high-level caching interfaces for memoization and
+result caching functionality.
+"""
+
+import atexit
+import functools
+import json
+import logging
+import pickle
+import shutil
+import threading
+import weakref
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from hashlib import sha256
+from os import PathLike
+from pathlib import Path
+from typing import cast, Generic, Protocol, TypedDict
+from typing_extensions import ParamSpec, TypeVar
+
+from filelock import FileLock
+
+import torch
+from torch._inductor.runtime.runtime_utils import cache_dir
+from torch._inductor.utils import clear_on_fresh_cache
+
+from . import config, implementations, locks
+
+
+logger = torch._logging.getArtifactLogger(__name__, "caching")
+
+# Type variable for function parameters
+_P = ParamSpec("_P")
+# Type variable for function return type
+_R = TypeVar("_R")
+# Type variable for encoded result type
+_EncodedR = TypeVar("_EncodedR")
+
+
+class CacheDumpEntry(TypedDict):
+    """A single cache entry in the dump format.
+
+    Attributes:
+        params: The encoded function parameters.
+        result: The encoded function result.
+    """
+
+    params: object
+    result: object
+
+
+class CacheDump(TypedDict):
+    """The structure of the memoizer cache dump file.
+
+    Cache entries are organized by sub_key in the collections field.
+    Memoizers without a sub_key use None as their key in collections.
+
+    Multiple Memoizer instances with different sub_keys can coexist in the same file.
+
+    Attributes:
+        collections: Dictionary mapping sub_keys (or None for root entries)
+            to their cache entries.
+        cache_size: The total number of cache entries across all collections.
+    """
+
+    collections: dict[str | None, dict[str, CacheDumpEntry]]
+    cache_size: int
+
+
+@dataclass
+class CacheEntry:
+    """A cache entry containing encoded parameters and result.
+
+    This dataclass stores the encoded form of function parameters and
+    the (possibly encoded) result for human-readable cache dumps.
+
+    Attributes:
+        encoded_params: The encoded function parameters used for debugging/inspection.
+        encoded_result: The (possibly encoded) result of the function call.
+    """
+
+    encoded_params: object
+    encoded_result: object
+
+
+@dataclass
+class InterimResult(Generic[_R]):
+    """Wrapper for an available interim result from a DeferredRecording.
+
+    When get_interim_result() returns an InterimResult, it means the
+    make_interim_result callable was provided and has been invoked.
+    The wrapped value is the result of that invocation.
+
+    This allows distinguishing between:
+    - No interim result configured: get_interim_result() returns None
+    - Interim result available: get_interim_result() returns InterimResult(value)
+      where value may itself be None if the callable returns None.
+
+    Attributes:
+        value: The value returned by calling make_interim_result().
+               May legitimately be None if the callable returns None.
+    """
+
+    value: _R
+
+
+@dataclass
+class DeferredRecording(Generic[_R, _EncodedR]):
+    """Signals that recording should happen at a later time.
+
+    When returned from a custom_result_encoder, the memoizer will:
+    1. Skip immediate caching
+    2. Register a completion callback on this object
+
+    The encoder is responsible for calling `finalize(encoded_result)`
+    when the actual result is ready to be cached. This is useful when
+    the function returns a value whose final encoded form is not yet
+    available (e.g., an object containing a pending computation),
+    allowing the expensive work to remain hidden while still enabling caching.
+
+    This class handles the race condition where the computation might complete
+    before the memoizer has registered its callback. If finalize() is called
+    first, the result is stored and the callbacks are invoked when registered.
+
+    Multiple callbacks can be registered - they will all be invoked when
+    finalize() is called (or immediately if already completed). Callbacks are
+    guaranteed to be invoked in registration order.
+
+    Optionally, a `make_interim_result` callable can be provided to handle
+    subsequent calls while the deferred recording is pending. When a memoized
+    function is called again with the same parameters while a deferred recording
+    is still pending, the memoizer can use this callable to construct an
+    appropriate return value instead of re-executing the function. The callable
+    does not necessarily return the same object - it provides a way to derive
+    a suitable response based on the pending deferred recording's context.
+
+    This class is thread-safe: concurrent calls to finalize() and
+    register_callback() are properly synchronized.
+
+    WARNING: Callbacks are executed while holding the internal lock to ensure
+    ordering guarantees. Callbacks MUST NOT call back into this DeferredRecording
+    instance (e.g., calling finalize() or register_callback() from within a
+    callback) as this will cause a deadlock.
+
+    Example usage in an encoder:
+        def my_encoder(*args, **kwargs):
+            def encode(result: _R) -> DeferredRecording[_R, _EncodedR]:
+                deferred: DeferredRecording[_R, _EncodedR] = DeferredRecording(
+                        # Construct an appropriate return for subsequent calls while pending
+                        make_interim_result=lambda: result
+                    )
+
+                def on_complete():
+                    encoded_result: _EncodedR = _encode_result(result)
+                    deferred.finalize(encoded_result)
+
+                result.add_completion_callback(on_complete)
+                return deferred
+
+            return encode
+
+    Attributes:
+        _callbacks: List of callbacks to invoke when finalize() is called.
+                   Set to None after finalize() is called to indicate completion.
+        _encoded_result: The encoded result passed to finalize().
+        _lock: Lock for thread-safe access to mutable state.
+        make_interim_result: Optional callable that constructs a return value for subsequent
+                           calls while the deferred recording is pending. If set, the memoizer
+                           calls this instead of re-executing the function. The callable may
+                           return the same object or construct a new appropriate response.
+    """
+
+    _callbacks: list[Callable[[_EncodedR], None]] | None = field(
+        default_factory=list, repr=False
+    )
+    _encoded_result: _EncodedR | None = field(default=None, repr=False)
+    _lock: threading.Lock = field(default_factory=threading.Lock, repr=False)
+    make_interim_result: Callable[[], _R] | None = field(default=None, repr=False)
+
+    def finalize(self, encoded_result: _EncodedR) -> None:
+        """Finalize the deferred recording with the encoded result.
+
+        This method should be called by the encoder (typically via a completion
+        callback) when the actual result is ready to be cached.
+
+        All registered callbacks are invoked with the encoded result in
+        registration order. If no callbacks are registered yet, the result is
+        stored and callbacks will be invoked when registered.
+
+        This method is thread-safe.
+
+        WARNING: Callbacks are executed while holding the internal lock.
+        Do not call finalize() or register_callback() from within a callback,
+        as this will cause a deadlock.
+
+        Args:
+            encoded_result: The encoded result to cache.
+        """
+        with self._lock:
+            if self._callbacks is None:
+                raise RuntimeError(
+                    "finalize() called multiple times on DeferredRecording"
+                )
+
+            # Store the result and get callbacks to invoke
+            self._encoded_result = encoded_result
+            callbacks_to_invoke = self._callbacks
+            self._callbacks = None
+
+            # Execute callbacks inside the lock to ensure ordering:
+            # - All pre-finalize callbacks execute in registration order
+            # - Any callback registered during/after finalize will see
+            #   _callbacks=None and execute immediately (also under lock)
+            for callback in callbacks_to_invoke:
+                callback(encoded_result)
+
+    def register_callback(self, callback: Callable[[_EncodedR], None]) -> None:
+        """Register a completion callback.
+
+        This method is called by memoizers to register callbacks that
+        will insert cache entries. Multiple callbacks can be registered.
+
+        If finalize() has already been called (e.g., the computation was fast),
+        the callback is invoked immediately with the stored result.
+
+        This method is thread-safe. Callbacks are guaranteed to be invoked
+        in registration order.
+
+        WARNING: Callbacks are executed while holding the internal lock.
+        Do not call finalize() or register_callback() from within a callback,
+        as this will cause a deadlock.
+
+        Args:
+            callback: The function to call with the encoded result.
+        """
+        with self._lock:
+            if self._callbacks is None:
+                # Already finalized - invoke immediately while holding lock
+                # to maintain ordering with respect to other callbacks
+                # we need to cast here, we can't assert because the encoded
+                # result may genuinely be None
+                callback(cast(_EncodedR, self._encoded_result))
+            else:
+                self._callbacks.append(callback)
+
+    def get_interim_result(self) -> InterimResult[_R] | None:
+        """Try to get an interim result for use while deferred recording is pending.
+
+        When a memoized function is called again with the same parameters while
+        a deferred recording is still pending, the memoizer uses this method to
+        check if an interim result can be returned instead of re-executing the
+        function.
+
+        This method provides thread-safe access to the make_interim_result
+        callable. It atomically checks if the callable was provided and invokes
+        it if so.
+
+        This method is thread-safe.
+
+        Returns:
+            InterimResult wrapping the value if make_interim_result was provided,
+            or None if no interim result is available.
+
+        Note:
+            The callable is invoked each time this method is called, so if
+            make_interim_result returns a new object each time, callers will
+            receive different objects. To return the same object, the callable
+            should capture and return a reference to a single object.
+        """
+        with self._lock:
+            if self.make_interim_result is not None:
+                return InterimResult(self.make_interim_result())
+            return None
+
+
+class ResultEncoderFactory(Protocol[_P, _R, _EncodedR]):
+    """Protocol for custom result encoder factories.
+
+    A result encoder factory is a callable with the following signature:
+        factory(fn) -> params_to_encoder(*args, **kwargs) -> encoder(result) -> encoded
+
+    The three levels are:
+    1. factory(fn): Takes the underlying unwrapped function, returns a params_to_encoder
+    2. params_to_encoder(*args, **kwargs): Takes the memoized function's arguments,
+       returns an encoder function
+    3. encoder(result): Converts the result R -> _EncodedR (or DeferredRecording)
+
+    The `fn` parameter allows the encoder to call the underlying function without
+    triggering memoization, which is useful when the encoder needs to re-execute
+    the function (e.g., to get a fresh result for DeferredRecording).
+
+    Example:
+        def my_encoder_factory(fn: Callable) -> Callable:
+            def params_to_encoder(*args, **kwargs) -> Callable:
+                def encode(result: R) -> EncodedR:
+                    return {"encoded": result}
+                return encode
+            return params_to_encoder
+
+        @memoizer.record(custom_result_encoder=my_encoder_factory)
+        def compute(x: int) -> int:
+            return x * 2
+    """
+
+    def __call__(
+        self,
+        fn: Callable[_P, _R],
+    ) -> Callable[_P, Callable[[_R], _EncodedR | DeferredRecording[_R, _EncodedR]]]: ...
+
+
+# pyrefly: ignore [variance-mismatch]
+class ResultDecoderFactory(Protocol[_P, _R, _EncodedR]):
+    """Protocol for custom result decoder factories.
+
+    A result decoder factory is a callable with the following signature:
+        factory(fn) -> params_to_decoder(*args, **kwargs) -> decoder(encoded) -> result
+
+    The three levels are:
+    1. factory(fn): Takes the underlying unwrapped function, returns a params_to_decoder
+    2. params_to_decoder(*args, **kwargs): Takes the memoized function's arguments,
+       returns a decoder function
+    3. decoder(encoded): Converts the encoded result _EncodedR -> R
+
+    The `fn` parameter allows the decoder to call the underlying function without
+    triggering memoization, which is useful for fallback paths when the cached
+    result cannot be fully decoded.
+
+    Example:
+        def my_decoder_factory(fn: Callable) -> Callable:
+            def params_to_decoder(*args, **kwargs) -> Callable:
+                def decode(encoded_result: EncodedR) -> R:
+                    return encoded_result["value"]
+                return decode
+            return params_to_decoder
+
+        @memoizer.replay(custom_result_decoder=my_decoder_factory)
+        def compute(x: int) -> int:
+            return x * 2
+    """
+
+    def __call__(
+        self,
+        fn: Callable[_P, _R],
+    ) -> Callable[_P, Callable[[_EncodedR], _R]]: ...
+
+
+class _BaseMemoizer:
+    """Base class for memoization interfaces.
+
+    This class provides the common memoize method that orchestrates
+    record and replay functionality.
+    """
+
+    @staticmethod
+    def _make_key(
+        custom_params_encoder: Callable[..., object] | None,
+        *args: object,
+        **kwargs: object,
+    ) -> str:
+        """Generate a cache key from function parameters.
+
+        Args:
+            custom_params_encoder: Optional encoder to apply to function parameters.
+                                  If None, params are pickled directly.
+            *args: Positional arguments to encode.
+            **kwargs: Keyword arguments to encode.
+
+        Returns:
+            A 32-character hex string suitable for use as a cache key.
+        """
+        if custom_params_encoder is None:
+            # Pickle the parameters directly
+            pickled_params: bytes = pickle.dumps((args, kwargs))
+        else:
+            # Encode the parameters using the custom encoder
+            encoded_params = custom_params_encoder(*args, **kwargs)
+            # Pickle the encoded output
+            pickled_params = pickle.dumps(encoded_params)
+
+        # Hash the pickled bytes with SHA256
+        hash_obj = sha256(pickled_params)
+
+        # Get hex digest and truncate to 32 characters
+        return hash_obj.hexdigest()[:32]
+
+    def record(
+        self,
+        custom_params_encoder: Callable[_P, object] | None = None,
+        custom_result_encoder: ResultEncoderFactory[_P, _R, _EncodedR] | None = None,
+    ) -> Callable[[Callable[_P, _R]], Callable[_P, _R]]:
+        """Record a function call result. Must be implemented by subclasses.
+
+        This can be used standalone to cache results without necessarily
+        replaying them, for example for logging or analytics purposes.
+
+        See memoize() for a description of the arguments.
+        """
+        raise NotImplementedError
+
+    def replay(
+        self,
+        custom_params_encoder: Callable[_P, object] | None = None,
+        custom_result_decoder: ResultDecoderFactory[_P, _R, _EncodedR] | None = None,
+    ) -> Callable[[Callable[_P, _R]], Callable[_P, _R]]:
+        """Replay a cached function result. Must be implemented by subclasses.
+
+        This can be used standalone to use cached values that are otherwise
+        too hard to automatically record.
+
+        See memoize() for a description of the arguments.
+        """
+        raise NotImplementedError
+
+    def memoize(
+        self,
+        custom_params_encoder: Callable[_P, object] | None = None,
+        custom_result_encoder: ResultEncoderFactory[_P, _R, _EncodedR] | None = None,
+        custom_result_decoder: ResultDecoderFactory[_P, _R, _EncodedR] | None = None,
+    ) -> Callable[[Callable[_P, _R]], Callable[_P, _R]]:
+        """Memoize a function with record and replay functionality.
+
+        This is a decorator that attempts to replay cached results first.
+        If a cache miss occurs, it records the result by executing the wrapped function.
+
+        Args:
+            custom_params_encoder: Optional encoder for function parameters.
+                                  If None, parameters are pickled directly.
+            custom_result_encoder: Optional encoder factory for function results.
+                                  Takes function parameters and returns an encoder
+                                  function that converts R -> _EncodedR.
+            custom_result_decoder: Optional decoder factory for cached results.
+                                  Takes function parameters and returns a decoder
+                                  function that converts _EncodedR -> R.
+
+        Returns:
+            A decorator function that can be applied to functions.
+
+        Example:
+            @memoizer.memoize(
+                custom_params_encoder=my_param_encoder,
+                custom_result_encoder=my_result_encoder_factory,
+                custom_result_decoder=my_result_decoder_factory,
+            )
+            def expensive_function(x, y):
+                return x + y
+        """
+
+        def wrapper(fn: Callable[_P, _R]) -> Callable[_P, _R]:
+            """Wrap the function to enable memoization with replay and record.
+
+            Args:
+                fn: The function to wrap.
+
+            Returns:
+                A wrapped version of the function.
+            """
+            # If caching is disabled, return the original function unchanged
+            if not config.IS_CACHING_MODULE_ENABLED():
+                return fn
+
+            # Create decorated versions using record and replay
+            replay_fn = self.replay(
+                custom_params_encoder,
+                custom_result_decoder,
+            )(fn)
+            record_fn = self.record(
+                custom_params_encoder,
+                custom_result_encoder,
+            )(fn)
+
+            @functools.wraps(fn)
+            def inner(*args: _P.args, **kwargs: _P.kwargs) -> _R:
+                """Attempt to replay from cache, or record on cache miss.
+
+                Args:
+                    *args: Positional arguments to pass to the function.
+                    **kwargs: Keyword arguments to pass to the function.
+
+                Returns:
+                    The result from cache (if hit) or from executing the function (if miss).
+                """
+                # Try to replay first
+                try:
+                    return replay_fn(*args, **kwargs)
+                except KeyError:
+                    # Cache miss - record the result
+                    return record_fn(*args, **kwargs)
+
+            return inner
+
+        return wrapper
+
+
+class Memoizer(_BaseMemoizer):
+    """In-memory memoization interface for caching function results.
+
+    This class provides methods for recording, retrieving, and managing
+    cached function results in memory with custom encoding/decoding logic.
+
+    Cache entries are stored as CacheEntry objects containing both the encoded
+    parameters and the encoded result. This makes debugging easier since entries
+    can be inspected with full context about what inputs produced each result.
+
+    On memoizer destruction, the cache is automatically dumped to a shared JSON file
+    in the cache directory for debugging and inspection purposes. Multiple Memoizer
+    instances contribute to the same file additively.
+
+    Note: Use this over functools.cache when you need to support parameters
+    that functools.cache cannot handle, or when you need custom encoding/decoding
+    of results.
+    """
+
+    def __init__(self, sub_key: str | None = None) -> None:
+        """Initialize the Memoizer instance with an in-memory cache.
+
+        Args:
+            sub_key: Optional key for organizing cache entries in the JSON dump.
+                    If provided, cache entries are stored under cache_entries[sub_key].
+                    If None, cache entries are merged directly into root cache_entries.
+        """
+        self._cache: implementations._InMemoryCacheImpl[CacheEntry] = (
+            implementations._InMemoryCacheImpl()
+        )
+        # Optional sub_key for nested cache structure
+        self._sub_key: str | None = sub_key
+        # Track pending deferred recordings by cache key.
+        # Uses WeakValueDictionary to prevent memory leaks: if the DeferredRecording
+        # is no longer referenced elsewhere (e.g., the async computation was cancelled
+        # or failed and garbage collected), the entry is automatically removed.
+        self._pending_deferred: weakref.WeakValueDictionary[str, DeferredRecording] = (
+            weakref.WeakValueDictionary()
+        )
+        # Register atexit handler to dump cache on program exit
+        if config.IS_DUMP_MEMOIZER_CACHE_ENABLED():
+            atexit.register(self._dump_to_disk)
+        # Pre-populate cache from dump file if configured (with sub_key now set)
+        self._maybe_prepopulate_from_dump()
+        # Register with clear_on_fresh_cache for fresh_cache() integration
+        clear_on_fresh_cache(self)
+
+    def cache_clear(self) -> None:
+        """Clear the in-memory cache.
+
+        This method resets the in-memory cache to empty. It is called by
+        the fresh_cache() context manager via clear_on_fresh_cache registration.
+        """
+        self._cache._memory.clear()
+
+    def _finalize_deferred_recording(
+        self,
+        cache_key: str,
+        encoded_params: object,
+        encoded_result: _EncodedR,
+    ) -> None:
+        """Finalize a deferred recording by inserting the cache entry.
+
+        This method is called when a DeferredRecording's finalize() method is invoked.
+        It creates the cache entry and inserts it into the in-memory cache.
+
+        Args:
+            cache_key: The cache key for the entry.
+            encoded_params: The encoded function parameters.
+            encoded_result: The final encoded result to cache.
+        """
+        # Remove from pending tracking
+        self._pending_deferred.pop(cache_key, None)
+        # Insert into cache
+        cache_entry = CacheEntry(
+            encoded_params=encoded_params,
+            encoded_result=encoded_result,
+        )
+        self._cache.insert(cache_key, cache_entry)
+
+    @functools.cached_property
+    def _shared_cache_filepath(self) -> Path:
+        """Get the shared cache filepath for memoizer cache dumps.
+
+        Returns:
+            The path to the shared memoizer cache JSON file.
+        """
+        return Path(cache_dir()) / "memoizer_cache.json"
+
+    @functools.cached_property
+    def _shared_cache_lockfile(self) -> Path:
+        """Get the lock file path for the shared memoizer cache.
+
+        Returns:
+            The path to the lock file for the shared cache.
+        """
+        return Path(cache_dir()) / "memoizer_cache.lock"
+
+    def _read_dump_from_disk(self, filepath: Path | None = None) -> CacheDump | None:
+        """Read a cache dump from disk.
+
+        Attempts to read and parse a cache JSON file.
+
+        Args:
+            filepath: Path to the dump file to read. If None, uses the
+                     shared cache filepath (self._shared_cache_filepath).
+
+        Returns:
+            The cache dump if the file exists and is valid JSON, None otherwise.
+        """
+        target_path = filepath if filepath is not None else self._shared_cache_filepath
+        try:
+            with open(target_path) as f:
+                data = json.load(f)
+                return data
+        except FileNotFoundError:
+            return None
+        except json.JSONDecodeError:
+            return None
+
+    def _write_dump_to_disk(self, dump: CacheDump) -> None:
+        """Write the cache dump to disk.
+
+        Writes the provided dump to the shared cache JSON file and logs the result.
+
+        Args:
+            dump: The cache dump to write.
+        """
+        try:
+            with open(self._shared_cache_filepath, "w") as f:
+                json.dump(dump, f, indent=2)
+
+            # Log the filepath
+            if self._sub_key:
+                logger.log(
+                    logging.INFO,
+                    "Memoizer cache (sub_key=%s) dumped to: %s",
+                    self._sub_key,
+                    self._shared_cache_filepath,
+                )
+            else:
+                logger.log(
+                    logging.INFO,
+                    "Memoizer cache dumped to: %s",
+                    self._shared_cache_filepath,
+                )
+        except Exception as e:
+            # If dumping fails, just log it and don't crash the program
+            logger.log(
+                logging.WARNING,
+                "Warning: Failed to dump memoizer cache: %s",
+                e,
+            )
+
+    def _prepare_dump(self, existing_dump: CacheDump | None) -> CacheDump:
+        """Prepare a cache dump from the current Memoizer state.
+
+        Takes the existing dump (if any) and merges it with the current
+        in-memory cache entries.
+
+        Args:
+            existing_dump: The existing dump to merge with, or None if starting fresh.
+
+        Returns:
+            The prepared cache dump ready to be written to disk.
+        """
+        # Start with existing data or empty structure
+        if existing_dump is not None:
+            dump = existing_dump
+        else:
+            dump: CacheDump = {"collections": {}, "cache_size": 0}
+
+        # Ensure collections field exists
+        if "collections" not in dump:
+            dump["collections"] = {}
+
+        # JSON serializes None keys as "null" string, so we need to handle both cases
+        # When reading from JSON, the key will be "null" (string), not None
+        lookup_key = "null" if self._sub_key is None else self._sub_key
+
+        # Get existing entries for this sub_key to merge with
+        existing_entries = dump["collections"].get(lookup_key, {})
+
+        # Format cache entries as {"params": ..., "result": ...}
+        formatted_cache: dict[str, CacheDumpEntry] = dict(existing_entries)
+        for key, value in self._cache._memory.items():
+            entry = value
+            formatted_cache[key] = CacheDumpEntry(
+                params=entry.encoded_params,
+                result=entry.encoded_result,
+            )
+
+        # Store under sub_key in collections (use lookup_key to maintain JSON compatibility)
+        dump["collections"][lookup_key] = formatted_cache
+
+        # Calculate total cache size
+        total_size = sum(len(collection) for collection in dump["collections"].values())
+        dump["cache_size"] = total_size
+
+        return dump
+
+    def _dump_to_disk(self) -> None:
+        """Dump the in-memory cache to a shared JSON file.
+
+        This method is automatically called on program exit via atexit.
+        It reads any existing cache data, merges it with this instance's cache,
+        and writes the combined result back. Multiple Memoizer instances
+        contribute to the same file additively.
+
+        If self._sub_key is set and non-empty, cache entries are stored under
+        cache_entries[sub_key]. Otherwise, they're merged directly into cache_entries.
+
+        Cache entries are formatted as {"params": <encoded_params>, "result": <encoded_result>}
+        for better human readability.
+
+        The filepath where the cache was dumped is logged.
+        """
+        # Skip if cache is empty
+        if not self._cache._memory:
+            return
+
+        # Ensure parent directory exists
+        self._shared_cache_filepath.parent.mkdir(parents=True, exist_ok=True)
+
+        # Acquire file lock to ensure thread/process safety
+        flock = FileLock(str(self._shared_cache_lockfile))
+        with locks._acquire_flock_with_timeout(flock):
+            existing_dump = self._read_dump_from_disk()
+            dump = self._prepare_dump(existing_dump)
+            self._write_dump_to_disk(dump)
+
+    def _maybe_prepopulate_from_dump(self) -> None:
+        """Pre-populate cache entries from a dump file if configured.
+
+        Checks the CACHE_DUMP_FILE_PATH config option for a path to a JSON dump file
+        produced by IS_DUMP_MEMOIZER_CACHE_ENABLED. If a valid path is provided and
+        the file exists, loads cache entries from it into the in-memory cache.
+
+        For Memoizer instances without a sub_key, loads entries from the root cache_entries.
+        For Memoizer instances with a sub_key, loads entries from cache_entries[sub_key].
+
+        This method is called during __init__ to pre-populate the cache.
+        """
+        dump_file_path = config.CACHE_DUMP_FILE_PATH()
+
+        # Skip if no dump file configured
+        if not dump_file_path:
+            return
+
+        # Read the dump file using the helper method
+        dump = self._read_dump_from_disk(Path(dump_file_path))
+        if dump is None:
+            return
+
+        # Extract entries to load from the dump
+        entries_to_load = self._extract_entries_from_dump(dump)
+        if not entries_to_load:
+            return
+
+        # Populate the cache
+        self._populate_cache_from_entries(entries_to_load)
+
+        # Log the result
+        if self._sub_key:
+            logger.log(
+                logging.INFO,
+                "Loaded %d cache entries from %s (sub_key=%s)",
+                len(entries_to_load),
+                dump_file_path,
+                self._sub_key,
+            )
+        else:
+            logger.log(
+                logging.INFO,
+                "Loaded %d cache entries from %s",
+                len(entries_to_load),
+                dump_file_path,
+            )
+
+    def _extract_entries_from_dump(self, dump: CacheDump) -> dict[str, CacheDumpEntry]:
+        """Extract cache entries from a dump based on sub_key.
+
+        Args:
+            dump: The cache dump to extract entries from.
+
+        Returns:
+            Dictionary of cache entries to load.
+        """
+        collections = dump.get("collections", {})
+        # JSON serializes None keys as "null" string, so we need to handle both cases
+        # When reading from JSON, the key will be "null" (string), not None
+        lookup_key = "null" if self._sub_key is None else self._sub_key
+        return collections.get(lookup_key, {})
+
+    def _populate_cache_from_entries(self, entries: dict[str, CacheDumpEntry]) -> None:
+        """Populate the in-memory cache from dump entries.
+
+        Args:
+            entries: Dictionary of cache entries to load.
+        """
+        for key, entry in entries.items():
+            cache_entry = CacheEntry(
+                encoded_params=entry["params"],
+                encoded_result=entry["result"],
+            )
+            self._cache.insert(key, cache_entry)
+
+    def record(
+        self,
+        custom_params_encoder: Callable[_P, object] | None = None,
+        custom_result_encoder: ResultEncoderFactory[_P, _R, _EncodedR] | None = None,
+    ) -> Callable[[Callable[_P, _R]], Callable[_P, _R]]:
+        """Record a function call result with custom encoding.
+
+        This is a decorator that wraps a function to enable memoization
+        with custom encoding/decoding logic.
+
+        Args:
+            custom_params_encoder: Optional encoder for function parameters.
+                                  If None, parameters are pickled directly.
+            custom_result_encoder: Optional encoder factory for function results.
+                                  First receives the underlying function, then takes
+                                  function parameters and returns an encoder function
+                                  that converts R -> _EncodedR.
+
+        Returns:
+            A decorator function that can be applied to functions.
+
+        Example:
+            @memoizer.record(
+                custom_params_encoder=my_param_encoder,
+                custom_result_encoder=my_result_encoder_factory,
+            )
+            def expensive_function(x, y):
+                return x + y
+        """
+
+        def wrapper(fn: Callable[_P, _R]) -> Callable[_P, _R]:
+            """Wrap the function to enable memoization.
+
+            Args:
+                fn: The function to wrap.
+
+            Returns:
+                A wrapped version of the function.
+            """
+            # If caching is disabled, return the original function unchanged
+            if not config.IS_CACHING_MODULE_ENABLED():
+                return fn
+
+            @functools.wraps(fn)
+            def inner(*args: _P.args, **kwargs: _P.kwargs) -> _R:
+                """Call the original function and cache the result.
+
+                Args:
+                    *args: Positional arguments to pass to the function.
+                    **kwargs: Keyword arguments to pass to the function.
+
+                Returns:
+                    The result of calling the original function.
+                """
+                # Call the function to compute the result
+                result = fn(*args, **kwargs)
+
+                # Generate cache key from parameters
+                cache_key = self._make_key(custom_params_encoder, *args, **kwargs)
+
+                # Encode params for human-readable dump
+                if custom_params_encoder is not None:
+                    encoded_params = custom_params_encoder(*args, **kwargs)
+                else:
+                    encoded_params = {
+                        "args": args,
+                        "kwargs": kwargs,
+                    }
+
+                # Encode the result if encoder is provided
+                if custom_result_encoder is not None:
+                    # Get the encoder function by calling the factory factory:
+                    # 1. First call with fn to get params_to_encoder factory
+                    # 2. Second call with params to get the encoder function
+                    params_to_encoder = custom_result_encoder(fn)
+                    encoder_fn = params_to_encoder(*args, **kwargs)
+                    encoded_result = encoder_fn(result)
+                else:
+                    encoded_result = result
+
+                # Check for deferred recording
+                if isinstance(encoded_result, DeferredRecording):
+                    # Track the pending deferred recording
+                    self._pending_deferred[cache_key] = encoded_result
+                    # Register the callback - handles race condition if encoded_result already completed
+                    encoded_result.register_callback(
+                        functools.partial(
+                            self._finalize_deferred_recording,
+                            cache_key,
+                            encoded_params,
+                        )
+                    )
+                    return result  # Return without caching
+
+                # Store CacheEntry in cache
+                cache_entry = CacheEntry(
+                    encoded_params=encoded_params,
+                    encoded_result=encoded_result,
+                )
+                self._cache.insert(cache_key, cache_entry)
+
+                # Return the original result (not the encoded version)
+                return result
+
+            return inner
+
+        return wrapper
+
+    def replay(
+        self,
+        custom_params_encoder: Callable[_P, object] | None = None,
+        custom_result_decoder: ResultDecoderFactory[_P, _R, _EncodedR] | None = None,
+    ) -> Callable[[Callable[_P, _R]], Callable[_P, _R]]:
+        """Replay a cached function result without executing the function.
+
+        This is a decorator that retrieves cached results instead of executing
+        the function. Raises an exception if no cached result exists.
+
+        Args:
+            custom_params_encoder: Optional encoder for function parameters.
+                                  If None, parameters are pickled directly.
+            custom_result_decoder: Optional decoder factory for cached results.
+                                  First receives the underlying function, then takes
+                                  function parameters and returns a decoder function
+                                  that converts _EncodedR -> R.
+
+        Returns:
+            A decorator function that can be applied to functions.
+
+        Example:
+            @memoizer.replay(
+                custom_params_encoder=my_param_encoder,
+                custom_result_decoder=my_result_decoder_factory,
+            )
+            def expensive_function(x, y):
+                return x + y
+        """
+
+        def wrapper(fn: Callable[_P, _R]) -> Callable[_P, _R]:
+            """Wrap the function to retrieve from cache.
+
+            Args:
+                fn: The function to wrap (not actually called).
+
+            Returns:
+                A wrapped version of the function.
+            """
+            # If caching is disabled, always raise KeyError (cache miss)
+            if not config.IS_CACHING_MODULE_ENABLED():
+
+                @functools.wraps(fn)
+                def always_miss(*args: _P.args, **kwargs: _P.kwargs) -> _R:
+                    raise KeyError("Caching is disabled")
+
+                return always_miss
+
+            @functools.wraps(fn)
+            def inner(*args: _P.args, **kwargs: _P.kwargs) -> _R:
+                """Retrieve the cached result without calling the function.
+
+                Args:
+                    *args: Positional arguments to generate the cache key.
+                    **kwargs: Keyword arguments to generate the cache key.
+
+                Returns:
+                    The cached result (decoded if decoder is provided).
+
+                Raises:
+                    KeyError: If no cached result exists for the given parameters.
+                """
+                # Generate cache key from parameters
+                cache_key = self._make_key(custom_params_encoder, *args, **kwargs)
+
+                # Check if result is cached
+                cached_hit = self._cache.get(cache_key)
+                if cached_hit is not None:
+                    # Extract the cached value
+                    cache_entry = cached_hit.value
+
+                    # Decode and return the cached result
+                    if custom_result_decoder is not None:
+                        # Get the decoder function by calling the factory factory:
+                        # 1. First call with fn to get params_to_decoder factory
+                        # 2. Second call with params to get the decoder function
+                        params_to_decoder = custom_result_decoder(fn)
+                        decoder_fn = params_to_decoder(*args, **kwargs)
+                        return decoder_fn(cast(_EncodedR, cache_entry.encoded_result))
+                    return cast(_R, cache_entry.encoded_result)
+
+                # Check for pending deferred recording with interim result
+                pending = self._pending_deferred.get(cache_key)
+                if pending is not None:
+                    interim = pending.get_interim_result()
+                    if interim is not None:
+                        return interim.value
+
+                raise KeyError(f"No cached result found for key: {cache_key}")
+
+            return inner
+
+        return wrapper
+
+
+class PersistentMemoizer(_BaseMemoizer):
+    """Persistent memoization interface for caching function results to disk.
+
+    This class provides methods for recording, retrieving, and managing
+    cached function results using a two-level cache strategy:
+    1. In-memory cache (fast, checked first) - via Memoizer instance
+    2. On-disk cache (persistent, checked on memory miss)
+
+    Results are persisted across process restarts.
+
+    On program exit, the in-memory cache entries are automatically dumped to
+    the shared JSON file. If sub_dir is non-empty, entries are stored under
+    a nested structure based on the sub_dir. If sub_dir is empty, entries are
+    merged directly into the root cache_entries.
+
+    Note: Use this over functools.cache when you need to support parameters
+    that functools.cache cannot handle, custom result encoding and/or decoding,
+    or when you need disk caching to persist results across program boundaries.
+    """
+
+    def __init__(self, sub_dir: PathLike[str] | None = None) -> None:
+        """Initialize the PersistentMemoizer with two-level caching.
+
+        Args:
+            sub_dir: Optional subdirectory within the cache directory for
+                    organizing cached results. Defaults to empty string if not specified.
+                    If non-empty, cache entries will be stored under cache_entries[sub_dir].
+                    If empty, cache entries are merged into root cache_entries.
+        """
+        # Use a Memoizer instance for in-memory caching
+        self._memoizer: Memoizer = Memoizer(sub_key=str(sub_dir) if sub_dir else None)
+        # Store on-disk cache as a separate attribute
+        self._disk_cache: implementations._OnDiskCacheImpl = (
+            implementations._OnDiskCacheImpl(sub_dir=sub_dir)
+        )
+        # Register with clear_on_fresh_cache for fresh_cache() integration
+        clear_on_fresh_cache(self)
+
+    def cache_clear(self) -> None:
+        """Clear the on-disk cache.
+
+        This method removes the on-disk cache directory. The in-memory cache
+        is cleared automatically by the underlying Memoizer instance, which
+        registers itself with clear_on_fresh_cache.
+
+        This is called by the fresh_cache() context manager via clear_on_fresh_cache
+        registration.
+        """
+        # Clear on-disk cache by removing the cache directory
+        # Note: in-memory cache is cleared by the Memoizer's own cache_clear
+        if self._disk_cache._cache_dir.exists():
+            shutil.rmtree(self._disk_cache._cache_dir, ignore_errors=True)
+
+    def record(
+        self,
+        custom_params_encoder: Callable[_P, object] | None = None,
+        custom_result_encoder: ResultEncoderFactory[_P, _R, _EncodedR] | None = None,
+    ) -> Callable[[Callable[_P, _R]], Callable[_P, _R]]:
+        """Record a function call result with custom encoding to both caches.
+
+        This is a decorator that wraps a function to enable memoization
+        with custom encoding/decoding logic. Results are stored in both
+        the in-memory cache and the on-disk cache.
+
+        This method delegates to the underlying Memoizer for memory caching,
+        then adds disk persistence. For deferred recordings, it registers
+        an additional callback to persist to disk when the recording completes.
+
+        Args:
+            custom_params_encoder: Optional encoder for function parameters.
+                                  If None, parameters are pickled directly.
+            custom_result_encoder: Optional encoder factory for function results.
+                                  First receives the underlying function, then takes
+                                  function parameters and returns an encoder function
+                                  that converts R -> _EncodedR.
+
+        Returns:
+            A decorator function that can be applied to functions.
+
+        Example:
+            @persistent_memoizer.record(
+                custom_params_encoder=my_param_encoder,
+                custom_result_encoder=my_result_encoder_factory,
+            )
+            def expensive_function(x, y):
+                return x + y
+        """
+
+        def wrapper(fn: Callable[_P, _R]) -> Callable[_P, _R]:
+            """Wrap the function to enable memoization.
+
+            Args:
+                fn: The function to wrap.
+
+            Returns:
+                A wrapped version of the function.
+            """
+            # If caching is disabled, return the original function unchanged
+            if not config.IS_CACHING_MODULE_ENABLED():
+                return fn
+
+            # Get the memory-cached version from the memoizer
+            memory_record_fn = self._memoizer.record(
+                custom_params_encoder, custom_result_encoder
+            )(fn)
+
+            @functools.wraps(fn)
+            def inner(*args: _P.args, **kwargs: _P.kwargs) -> _R:
+                """Call the original function and cache the result in both caches.
+
+                Args:
+                    *args: Positional arguments to pass to the function.
+                    **kwargs: Keyword arguments to pass to the function.
+
+                Returns:
+                    The result of calling the original function.
+                """
+                # Call the memory-cached version (which calls fn and caches in memory)
+                result = memory_record_fn(*args, **kwargs)
+
+                # Also store in disk cache
+                cache_key = self._make_key(custom_params_encoder, *args, **kwargs)
+
+                # Check if this was a deferred recording
+                pending = self._memoizer._pending_deferred.get(cache_key)
+                if pending is not None:
+                    # Deferred recording - register our own callback for disk persistence
+                    # By the time our callback runs, Memoizer's callback will have
+                    # already inserted the cache entry, so we just read and persist it
+                    pending.register_callback(
+                        functools.partial(self._persist_to_disk, cache_key)
+                    )
+                    return result
+
+                # Persist the cache entry to disk
+                self._persist_to_disk(cache_key)
+
+                return result
+
+            return inner
+
+        return wrapper
+
+    def _persist_to_disk(
+        self,
+        cache_key: str,
+        _callback_result: object = None,
+    ) -> None:
+        """Persist a cache entry to disk.
+
+        This method handles disk persistence for both immediate and deferred recordings:
+        - For immediate recordings: called directly after memory caching
+        - For deferred recordings: registered as a callback that runs after the
+          Memoizer's callback has inserted the entry into memory
+
+        Args:
+            cache_key: The cache key for the entry.
+            _callback_result: Unused. When called as a callback from DeferredRecording,
+                             this receives the encoded result, but we ignore it and
+                             read the full CacheEntry from the Memoizer's cache instead.
+        """
+        # Always read from memory cache - Memoizer's callback has already inserted it
+        cached_hit = self._memoizer._cache.get(cache_key)
+        assert cached_hit, "Cache entry must exist in memory cache"
+        cache_entry = cached_hit.value
+
+        # Store the full CacheEntry in disk cache for easier debugging
+        pickled_entry: bytes = pickle.dumps(cache_entry)
+        self._disk_cache.insert(cache_key, pickled_entry)
+
+    def replay(
+        self,
+        custom_params_encoder: Callable[_P, object] | None = None,
+        custom_result_decoder: ResultDecoderFactory[_P, _R, _EncodedR] | None = None,
+    ) -> Callable[[Callable[_P, _R]], Callable[_P, _R]]:
+        """Replay a cached function result without executing the function.
+
+        This is a decorator that retrieves cached results using a two-level
+        cache strategy. It checks the in-memory cache first (fast), then
+        falls back to the on-disk cache. If found on disk, the result is
+        cached in memory for future access.
+
+        Args:
+            custom_params_encoder: Optional encoder for function parameters.
+                                  If None, parameters are pickled directly.
+            custom_result_decoder: Optional decoder factory for cached results.
+                                  First receives the underlying function, then takes
+                                  function parameters and returns a decoder function
+                                  that converts _EncodedR -> R.
+
+        Returns:
+            A decorator function that can be applied to functions.
+
+        Example:
+            @persistent_memoizer.replay(
+                custom_params_encoder=my_param_encoder,
+                custom_result_decoder=my_result_decoder_factory,
+            )
+            def expensive_function(x, y):
+                return x + y
+        """
+
+        def wrapper(fn: Callable[_P, _R]) -> Callable[_P, _R]:
+            """Wrap the function to retrieve from cache.
+
+            Args:
+                fn: The function to wrap (not actually called).
+
+            Returns:
+                A wrapped version of the function.
+            """
+            # If caching is disabled, always raise KeyError (cache miss)
+            if not config.IS_CACHING_MODULE_ENABLED():
+
+                @functools.wraps(fn)
+                def always_miss(*args: _P.args, **kwargs: _P.kwargs) -> _R:
+                    raise KeyError("Caching is disabled")
+
+                return always_miss
+
+            # Get the memory replay function
+            memory_replay_fn = self._memoizer.replay(
+                custom_params_encoder, custom_result_decoder
+            )(fn)
+
+            @functools.wraps(fn)
+            def inner(*args: _P.args, **kwargs: _P.kwargs) -> _R:
+                """Retrieve the cached result without calling the function.
+
+                Checks memory cache first, then disk cache. Populates memory
+                cache from disk on a disk hit.
+
+                Args:
+                    *args: Positional arguments to generate the cache key.
+                    **kwargs: Keyword arguments to generate the cache key.
+
+                Returns:
+                    The cached result (decoded if decoder is provided).
+
+                Raises:
+                    KeyError: If no cached result exists for the given parameters.
+                """
+                # Try memory cache first via memoizer
+                try:
+                    return memory_replay_fn(*args, **kwargs)
+                except KeyError:
+                    pass  # Memory miss, check disk
+
+                # Memory miss - check disk cache
+                cache_key = self._make_key(custom_params_encoder, *args, **kwargs)
+                disk_hit = self._disk_cache.get(cache_key)
+                if disk_hit is not None:
+                    # Disk cache hit - unpickle the CacheEntry
+                    pickled_value = disk_hit.value
+                    cache_entry = cast(CacheEntry, pickle.loads(pickled_value))
+
+                    # Populate memory cache for future access
+                    self._memoizer._cache.insert(cache_key, cache_entry)
+
+                    # Decode and return
+                    if custom_result_decoder is not None:
+                        # Get the decoder function by calling the factory factory:
+                        # 1. First call with fn to get params_to_decoder factory
+                        # 2. Second call with params to get the decoder function
+                        params_to_decoder = custom_result_decoder(fn)
+                        decoder_fn = params_to_decoder(*args, **kwargs)
+                        return decoder_fn(cast(_EncodedR, cache_entry.encoded_result))
+                    return cast(_R, cache_entry.encoded_result)
+
+                # Complete miss
+                raise KeyError(f"No cached result found for key: {cache_key}")
+
+            return inner
+
+        return wrapper

outputs/audit_venv/lib/python3.11/site-packages/torch/_inductor/runtime/caching/locks.py

@@ -0,0 +1,207 @@
+"""Lock acquisition utilities
+
+This module provides safe and unsafe lock acquisition functions for both threading.Lock
+and FileLock objects, with configurable timeout behaviors. It supports three timeout modes:
+blocking (infinite wait), non-blocking (immediate), and blocking with timeout (finite wait).
+
+The module offers both context manager and manual acquisition patterns:
+- Safe acquisition: Uses context managers that automatically handle lock release
+- Unsafe acquisition: Manual acquisition that requires explicit release by the caller
+"""
+
+from __future__ import annotations
+
+from contextlib import _GeneratorContextManager, contextmanager, ExitStack
+from typing import TYPE_CHECKING, TypeAlias
+from typing_extensions import Protocol
+
+from filelock import BaseFileLock, Timeout
+
+from . import exceptions
+
+
+if TYPE_CHECKING:
+    from collections.abc import Generator
+    from threading import Lock
+
+    from .implementations import _CacheImpl
+
+
+_LockContextManager: TypeAlias = _GeneratorContextManager[None, None, None]
+
+
+class _LockProtocol(Protocol):  # noqa: PYI046
+    def __call__(self, timeout: float | None = None) -> _LockContextManager: ...
+
+
+# Infinite timeout - blocks indefinitely until lock is acquired.
+_BLOCKING: float = -1
+# No timeout - returns immediately if lock cannot be acquired.
+_NON_BLOCKING: float = 0
+# Finite timeout - blocks for a specified duration before raising a timeout error.
+_BLOCKING_WITH_TIMEOUT: float = 60.0
+# Default timeout for lock acquisition.
+_DEFAULT_TIMEOUT: float = _BLOCKING_WITH_TIMEOUT
+
+
+@contextmanager
+def _acquire_lock_with_timeout(
+    lock: Lock,
+    timeout: float | None = None,
+) -> Generator[None, None, None]:
+    """Context manager that safely acquires a threading.Lock with timeout and automatically releases it.
+
+    This function provides a safe way to acquire a lock with timeout support, ensuring
+    the lock is always released even if an exception occurs during execution.
+
+    Args:
+        lock: The threading.Lock object to acquire
+        timeout: Timeout in seconds. If None, uses _DEFAULT_TIMEOUT.
+                - Use _BLOCKING (-1.0) for infinite wait
+                - Use _NON_BLOCKING (0.0) for immediate return
+                - Use positive value for finite timeout
+
+    Yields:
+        None: Yields control to the caller while holding the lock
+
+    Raises:
+        LockTimeoutError: If the lock cannot be acquired within the timeout period
+
+    Example:
+        with _acquire_lock_with_timeout(my_lock, timeout=30.0):
+            # Critical section - lock is held
+            perform_critical_operation()
+        # Lock is automatically released here
+    """
+    _unsafe_acquire_lock_with_timeout(lock, timeout=timeout)
+
+    try:
+        yield
+    finally:
+        lock.release()
+
+
+def _unsafe_acquire_lock_with_timeout(lock: Lock, timeout: float | None = None) -> None:
+    """Acquire a threading.Lock with timeout without automatic release (unsafe).
+
+    This function acquires a lock with timeout support but does NOT automatically
+    release it. The caller is responsible for releasing the lock explicitly.
+    Use this only when you need manual control over lock lifetime.
+
+    Args:
+        lock: The threading.Lock object to acquire
+        timeout: Timeout in seconds. If None, uses _DEFAULT_TIMEOUT.
+                - Use _BLOCKING (-1.0) for infinite wait
+                - Use _NON_BLOCKING (0.0) for immediate return
+                - Use positive value for finite timeout
+
+    Raises:
+        LockTimeoutError: If the lock cannot be acquired within the timeout period
+
+    Warning:
+        This is an "unsafe" function because it does not automatically release
+        the lock. Always call lock.release() when done, preferably in a try/finally
+        block or use the safe _acquire_lock_with_timeout context manager instead.
+
+    Example:
+        lock = Lock()
+        try:
+            _unsafe_acquire_lock_with_timeout(lock, timeout=30.0)
+            # Critical section - lock is held
+            perform_critical_operation()
+        finally:
+            lock.release()  # Must manually release!
+    """
+    _timeout: float = timeout if timeout is not None else _DEFAULT_TIMEOUT
+    if not lock.acquire(timeout=_timeout):
+        raise exceptions.LockTimeoutError(lock, _timeout)
+
+
+@contextmanager
+def _acquire_flock_with_timeout(
+    flock: BaseFileLock,
+    timeout: float | None = None,
+) -> Generator[None, None, None]:
+    """Context manager that safely acquires a FileLock with timeout and automatically releases it.
+
+    This function provides a safe way to acquire a file lock with timeout support, ensuring
+    the lock is always released even if an exception occurs during execution.
+
+    Args:
+        flock: The FileLock object to acquire
+        timeout: Timeout in seconds. If None, uses _DEFAULT_TIMEOUT.
+                - Use _BLOCKING (-1.0) for infinite wait
+                - Use _NON_BLOCKING (0.0) for immediate return
+                - Use positive value for finite timeout
+
+    Yields:
+        None: Yields control to the caller while holding the file lock
+
+    Raises:
+        FileLockTimeoutError: If the file lock cannot be acquired within the timeout period
+
+    Example:
+        flock = FileLock("/tmp/my_process.lock")
+        with _acquire_flock_with_timeout(flock, timeout=30.0):
+            # Critical section - file lock is held
+            perform_exclusive_file_operation()
+        # File lock is automatically released here
+    """
+    _unsafe_acquire_flock_with_timeout(flock, timeout=timeout)
+
+    try:
+        yield
+    finally:
+        flock.release()
+
+
+def _unsafe_acquire_flock_with_timeout(
+    flock: BaseFileLock,
+    timeout: float | None,
+) -> None:
+    """Acquire a FileLock with timeout without automatic release (unsafe).
+
+    This function acquires a file lock with timeout support but does NOT automatically
+    release it. The caller is responsible for releasing the lock explicitly.
+    Use this only when you need manual control over lock lifetime.
+
+    Args:
+        flock: The FileLock object to acquire
+        timeout: Timeout in seconds. If None, uses _DEFAULT_TIMEOUT.
+                - Use _BLOCKING (-1.0) for infinite wait
+                - Use _NON_BLOCKING (0.0) for immediate return
+                - Use positive value for finite timeout
+
+    Raises:
+        FileLockTimeoutError: If the file lock cannot be acquired within the timeout period
+
+    Warning:
+        This is an "unsafe" function because it does not automatically release
+        the lock. Always call flock.release() when done, preferably in a try/finally
+        block or use the safe _acquire_flock_with_timeout context manager instead.
+
+    Example:
+        flock = FileLock("/tmp/my_process.lock")
+        try:
+            _unsafe_acquire_flock_with_timeout(flock, timeout=30.0)
+            # Critical section - file lock is held
+            perform_exclusive_file_operation()
+        finally:
+            flock.release()  # Must manually release!
+    """
+    _timeout: float = timeout if timeout is not None else _DEFAULT_TIMEOUT
+    try:
+        _ = flock.acquire(timeout=_timeout)
+    except Timeout as err:
+        raise exceptions.FileLockTimeoutError(flock, _timeout) from err
+
+
+@contextmanager
+def _acquire_many_impl_locks_with_timeout(
+    *impls: _CacheImpl,
+    timeout: float | None = None,
+) -> Generator[None, None, None]:
+    with ExitStack() as stack:
+        for impl in impls:
+            stack.enter_context(impl.lock(timeout))
+        yield

outputs/audit_venv/lib/python3.11/site-packages/torch/_inductor/runtime/caching/memoizers.py

@@ -0,0 +1,7 @@
+from pathlib import Path
+
+from .interfaces import PersistentMemoizer
+
+
+# Memoizer for _should_pad in pad_mm
+should_pad_memoizer = PersistentMemoizer(sub_dir=Path("should_pad"))

outputs/audit_venv/lib/python3.11/site-packages/torch/_inductor/runtime/caching/utils.py

@@ -0,0 +1,66 @@
+"""Utility functions
+
+This module provides helper functions for LRU caching decorators used
+throughout the caching system.
+"""
+
+from collections.abc import Callable
+from functools import lru_cache, wraps
+from typing_extensions import ParamSpec, TypedDict, TypeVar
+
+from torch import Tensor
+
+
+# Type specification for function parameters
+P = ParamSpec("P")
+# Type variable for function return values
+R = TypeVar("R")
+
+
+def _lru_cache(fn: Callable[P, R]) -> Callable[P, R]:
+    """LRU cache decorator with TypeError fallback.
+
+    Provides LRU caching with a fallback mechanism that calls the original
+    function if caching fails due to unhashable arguments. Uses a cache
+    size of 64 with typed comparison.
+
+    Args:
+        fn: The function to be cached.
+
+    Returns:
+        A wrapper function that attempts caching with fallback to original function.
+    """
+    cached_fn = lru_cache(maxsize=64, typed=True)(fn)
+
+    @wraps(fn)
+    def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+        try:
+            return cached_fn(*args, **kwargs)
+        except TypeError:
+            return fn(*args, **kwargs)
+
+    return wrapper
+
+
+class EncodedTensor(TypedDict):
+    """TypedDict for encoded tensor metadata."""
+
+    shape: tuple[int, ...]
+    stride: tuple[int, ...]
+    dtype: str
+
+
+def _encode_tensor(t: Tensor) -> EncodedTensor:
+    """Encode a tensor's metadata into a JSON-serializable dict.
+
+    Args:
+        t: PyTorch tensor to encode
+
+    Returns:
+        Dict containing shape, stride, and dtype information
+    """
+    return EncodedTensor(
+        shape=tuple(t.shape),
+        stride=tuple(t.stride()),
+        dtype=str(t.dtype),
+    )

outputs/audit_venv/lib/python3.11/site-packages/torch/_inductor/runtime/compile_tasks.py

@@ -0,0 +1,150 @@
+from __future__ import annotations
+
+import functools
+import linecache
+import os
+import sys
+import time
+import warnings
+from pathlib import Path
+from types import ModuleType
+from typing import Any, TYPE_CHECKING
+
+from torch._utils_internal import log_triton_builds
+
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from torch._inductor.runtime.triton_heuristics import CachingAutotuner
+
+
+def _reload_python_module(
+    key: str, path: str, set_sys_modules: bool = True
+) -> ModuleType:
+    with open(path) as f:
+        try:
+            code = compile(f.read(), path, "exec", dont_inherit=True)
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to import {path}\n{type(e).__name__}: {e}"
+            ) from None
+        mod = ModuleType(f"{__name__}.{key}")
+        mod.__file__ = path
+        mod.key = key  # type: ignore[attr-defined]
+        exec(code, mod.__dict__, mod.__dict__)
+        if set_sys_modules:
+            sys.modules[mod.__name__] = mod
+        return mod
+
+
+@functools.cache
+def _set_triton_ptxas_path() -> None:
+    if os.environ.get("TRITON_PTXAS_PATH") is not None:
+        return
+    ptxas = Path(__file__).absolute().parents[2] / "bin" / "ptxas"
+    if not ptxas.exists():
+        return
+    if ptxas.is_file() and os.access(ptxas, os.X_OK):
+        os.environ["TRITON_PTXAS_PATH"] = str(ptxas)
+    else:
+        warnings.warn(f"{ptxas} exists but is not an executable")
+
+
+def _set_triton_libdevice_path() -> None:
+    """
+    Use the CUDA toolkit's libdevice instead of Triton's bundled version.
+    This ensures Triton's pow matches CUDA's powf for bitwise precision.
+    Gated by config.eager_numerics.use_pytorch_libdevice.
+    """
+    from torch._inductor import config
+
+    if not config.eager_numerics.use_pytorch_libdevice:
+        return
+
+    _set_triton_libdevice_path_impl()
+
+
+def _set_triton_libdevice_path_impl() -> None:
+    try:
+        from triton import knobs
+    except ImportError:
+        return
+
+    env_path = os.environ.get("TRITON_LIBDEVICE_PATH")
+    if env_path is not None:
+        knobs.nvidia.libdevice_path = env_path
+        return
+
+    if knobs.nvidia.libdevice_path is not None:
+        return
+
+    try:
+        from torch.utils.cpp_extension import CUDA_HOME
+
+        if CUDA_HOME is None:
+            warnings.warn(
+                "CUDA_HOME not set; using Triton's bundled libdevice which may "
+                "cause minor precision differences in pow operations. "
+                "To fix: set TRITON_LIBDEVICE_PATH to your CUDA toolkit's libdevice, "
+                "e.g., export TRITON_LIBDEVICE_PATH=/usr/local/cuda/nvvm/libdevice/libdevice.10.bc",
+                stacklevel=3,
+            )
+            return
+        libdevice = Path(CUDA_HOME) / "nvvm" / "libdevice" / "libdevice.10.bc"
+        if libdevice.is_file():
+            knobs.nvidia.libdevice_path = str(libdevice)
+            # Also set env var so subprocess compile workers inherit it
+            os.environ["TRITON_LIBDEVICE_PATH"] = str(libdevice)
+        else:
+            warnings.warn(
+                f"CUDA libdevice not found at {libdevice}; using Triton's bundled "
+                "libdevice which may cause minor precision differences in pow operations. "
+                "To fix: set TRITON_LIBDEVICE_PATH to your CUDA toolkit's libdevice, "
+                "e.g., export TRITON_LIBDEVICE_PATH=/usr/local/cuda/nvvm/libdevice/libdevice.10.bc",
+                stacklevel=3,
+            )
+    except ImportError:
+        warnings.warn(
+            "torch.utils.cpp_extension not available; using Triton's bundled "
+            "libdevice which may cause minor precision differences in pow operations. "
+            "To fix: set TRITON_LIBDEVICE_PATH to your CUDA toolkit's libdevice, "
+            "e.g., export TRITON_LIBDEVICE_PATH=/usr/local/cuda/nvvm/libdevice/libdevice.10.bc",
+            stacklevel=3,
+        )
+
+
+def _worker_compile_triton(
+    load_kernel: Callable[[], CachingAutotuner],
+    extra_env: dict[str, str],
+    extra_config: dict[str, Any],
+) -> tuple[CachingAutotuner, int]:
+    _set_triton_ptxas_path()
+    os.environ.update(extra_env)
+    # Set libdevice path if passed via env from main process
+    libdevice_path = extra_env.get("TRITON_LIBDEVICE_PATH")
+    if libdevice_path:
+        try:
+            from triton import knobs
+
+            knobs.nvidia.libdevice_path = libdevice_path
+        except ImportError:
+            pass
+    from torch._inductor import config
+
+    with config.patch(extra_config):
+        fail = None
+        try:
+            start_ns = time.time_ns()
+            kernel = load_kernel()
+            kernel.precompile(warm_cache_only=True)
+            elapsed_ns = time.time_ns() - start_ns
+            kernel.prepare_for_pickle()
+            # We can release this memory in the compile subprocesses:
+            linecache.clearcache()
+            return kernel, elapsed_ns // 1000
+        except Exception as e:
+            fail = str(e)
+            raise
+        finally:
+            log_triton_builds(fail=fail)

outputs/audit_venv/lib/python3.11/site-packages/torch/_inductor/runtime/coordinate_descent_tuner.py

@@ -0,0 +1,412 @@
+# mypy: allow-untyped-defs
+import copy
+import itertools
+import logging
+from collections.abc import Callable
+from typing import TYPE_CHECKING
+
+from torch.utils._ordered_set import OrderedSet
+
+from ..utils import get_max_numwarps
+from .hints import TRITON_MAX_BLOCK
+from .runtime_utils import red_text, triton_config_to_hashable
+
+
+if TYPE_CHECKING:
+    from .triton_compat import triton
+
+
+log = logging.getLogger(__name__)
+
+
+def get_field(config, name):
+    if name == "num_warps":
+        return config.num_warps
+    elif name == "num_stages":
+        return config.num_stages
+    elif name == "waves_per_eu":
+        return config.kwargs.get(name, int(8 // config.num_warps))
+    else:
+        return config.kwargs.get(name, None)
+
+
+def set_field(config, name, value):
+    if name == "num_warps":
+        config.num_warps = value
+    elif name == "num_stages":
+        config.num_stages = value
+    else:
+        config.kwargs[name] = value
+
+
+class CoordescTuner:
+    """
+    The coordinate descent tuner. Tune one field/coordinate at a time.
+
+    TODO will it be necessary to tune multiple fields simultaneously.
+
+
+    TODO: what if both increasing and decreasing a field can improve perf.
+          i.e., there are multiple local optima..
+    """
+
+    def __init__(
+        self,
+        is_mm=False,
+        is_native_matmul=False,
+        is_mix_order_reduction=False,
+        name="unknown",
+        size_hints=None,
+        inductor_meta=None,
+        frozen_fields=None,
+    ):
+        self.is_mm = is_mm  # we will tune num_stages for mm
+
+        # Native matmul codegen assumes ZBLOCK=1 always.
+        # This is because 3d tl.dot is slow and so we want to tile y and x only.
+        # tl.dot also does not support size smaller than 16; we put this restriction.
+        self.is_native_matmul = is_native_matmul
+        assert not (self.is_mm and self.is_native_matmul)
+        self.is_mix_order_reduction = is_mix_order_reduction
+        self.cached_benchmark_results = {}
+        self.name = name
+        self.size_hints = size_hints
+        self.inductor_meta = inductor_meta or {}
+        self.frozen_fields: OrderedSet[str] = (
+            OrderedSet(frozen_fields) if frozen_fields is not None else OrderedSet()
+        )
+
+    def get_config_max(self, prefix: str) -> int:
+        max_block = TRITON_MAX_BLOCK[prefix.upper()]
+        size_hint = self.size_hints.get(prefix) if self.size_hints is not None else None
+        return min(max_block, size_hint) if size_hint is not None else max_block
+
+    def get_warpsmax(self):
+        # Avoid querying device directly if device properties are populated in inductor_meta
+        warp_size = self.inductor_meta.get("warp_size")
+        max_threads_per_block = self.inductor_meta.get("max_threads_per_block")
+        if warp_size and max_threads_per_block:
+            return max_threads_per_block // warp_size
+        else:
+            return get_max_numwarps()
+
+    def cache_benchmark_result(self, config, timing):
+        self.cached_benchmark_results[triton_config_to_hashable(config)] = timing
+
+    def lookup_in_cache(self, config):
+        return self.cached_benchmark_results.get(triton_config_to_hashable(config))
+
+    def call_func(self, func, config):
+        found = self.lookup_in_cache(config)
+        if found is not None:
+            log.debug("  CACHED")
+            return found
+        timing = func(config)
+        self.cache_benchmark_result(config, timing)
+        return timing
+
+    @property
+    def tunable_fields(self):
+        out = [
+            "XBLOCK",
+            "YBLOCK",
+            "ZBLOCK",
+            # NOTE: we should not tune R0_BLOCK for persistent reduction.
+            # We rely on the fact that persistent reduction's triton.Config
+            # does not have the R0_BLOCK field to guarantee that.
+            "R0_BLOCK",
+            "R1_BLOCK",
+            # the following 3 are for mm
+            "BLOCK_M",
+            "BLOCK_N",
+            "BLOCK_K",
+            "num_warps",
+        ]
+        if self.is_mm:
+            out.append("num_stages")
+        if self.inductor_meta.get("is_hip") is True:
+            out.append("waves_per_eu")
+        if self.is_native_matmul:
+            out.append("num_stages")
+            out.remove("ZBLOCK")  # ZBLOCK=1 always in native matmul
+
+        if self.is_mix_order_reduction:
+            # unlike TritonConfig.num_stages, this one is
+            # put in TritonConfig.kwargs["NUM_STAGES"] and is used to
+            # control the stage of pipelining of tl.range.
+            out.append("NUM_STAGES")
+
+        return [f for f in out if f not in self.frozen_fields]
+
+    def value_too_large(self, name: str, val: int) -> bool:
+        block_suffix = "BLOCK"
+        if name.endswith(block_suffix):
+            prefix = name.strip(block_suffix).lower()
+            return val > self.get_config_max(prefix)
+        if name == "num_warps":
+            return val > self.get_warpsmax()
+        if name == "waves_per_eu":
+            return val > 8
+
+        return False
+
+    def value_too_small(self, name: str, val: int) -> bool:
+        # In native matmul, block size should be >= 16 for tl.dot
+        if self.is_native_matmul:
+            if name in ["YBLOCK", "XBLOCK", "R0_BLOCK"]:
+                return val < 16
+
+        # Break if value becomes 0/neg
+        return val <= 0
+
+    def get_neighbour_values(self, name, orig_val, radius=None, include_self=False):
+        """
+        Get neighbour values in 'radius' steps. The original value is not
+        returned as it's own neighbour.
+        """
+        if radius is None:
+            radius = 1
+        if name == "NUM_STAGES":
+            # we see cases that
+            # NUM_STAGES=1 is better than NUM_STAGES=2
+            # while NUM_STAGES=1 is worse than NUM_STAGES=3
+            radius = max(radius, 2)
+
+        assert radius >= 1
+
+        def update(cur_val, inc=True):
+            if name in ["num_stages", "NUM_STAGES"]:
+                if inc:
+                    return cur_val + 1
+                else:
+                    return cur_val - 1
+            else:
+                if inc:
+                    return cur_val * 2
+                else:
+                    return cur_val // 2
+
+        out = []
+        # increment loop
+        cur_val = orig_val
+        for _ in range(radius):
+            cur_val = update(cur_val, True)
+            if self.value_too_large(name, cur_val):
+                break
+            out.append(cur_val)
+
+        # decrement loop
+        cur_val = orig_val
+        for _ in range(radius):
+            cur_val = update(cur_val, False)
+            if self.value_too_small(name, cur_val):
+                break
+            out.append(cur_val)
+
+        if include_self:
+            out.append(orig_val)
+        return out
+
+    @staticmethod
+    def has_improvement(baseline, test):
+        threshold = 0.001  # 0.1%
+        return test is not None and test < baseline * (1 - threshold)
+
+    def is_valid_config(self, config) -> bool:
+        if self.is_mix_order_reduction:
+            # Mix order reduction has an extra constraint that
+            # we should not tune XBLOCK beyond RSPLIT_SIZE
+            xblock = config.kwargs["XBLOCK"]
+            split_size = config.kwargs["RSPLIT_SIZE"]
+            return xblock <= split_size
+        return True
+
+    def check_all_tuning_directions(
+        self,
+        # pyrefly: ignore [missing-attribute]
+        func: Callable[["triton.Config"], float],
+        best_config,
+        best_timing,
+    ):
+        """
+        Check all directions. We only do this once the regular coordinate
+        descent tuning find no better choices any more.
+        We only have a few tunable fields, so this should be fine.
+        """
+        candidate_values_list = []
+        effective_fields = []
+        for field in self.tunable_fields:
+            old_value = get_field(best_config, field)
+            if old_value is None:
+                continue
+            radius = self.inductor_meta.get("coordinate_descent_search_radius", 1)
+            candidate_values = self.get_neighbour_values(
+                field,
+                old_value,
+                radius=radius,
+                include_self=True,
+            )
+            candidate_values_list.append(candidate_values)
+            effective_fields.append(field)
+
+        choices = itertools.product(*candidate_values_list)
+        improved = False
+        for choice in choices:
+            assert len(choice) == len(effective_fields)
+            candidate_config = copy.deepcopy(best_config)
+            for new_val, field in zip(choice, effective_fields):
+                set_field(candidate_config, field, new_val)
+            if not self.is_valid_config(candidate_config):
+                continue
+            cmp_res, candidate_timing = self.compare_config(
+                func, candidate_config, best_config, best_timing
+            )
+            if cmp_res:
+                improved = True
+                best_config = candidate_config
+                best_timing = candidate_timing
+
+        return improved, best_config, best_timing
+
+    def compare_config(self, func, candidate_config, best_config, best_timing):
+        """
+        Check if candidate_config is better than best_config.
+
+        Return a tuple of (compare_result, candidate_timing).
+        compare_result is true iff candidate_config is better.
+        """
+        log.debug("Try config %s", candidate_config)
+        try:
+            candidate_timing = self.call_func(func, candidate_config)
+        except Exception as e:
+            log.debug("Got exception %s", e)  # noqa: G200
+            return False, float("inf")
+
+        if self.has_improvement(best_timing, candidate_timing):
+            log.debug(
+                "Tune from %s %f -> %s %f",
+                best_config,
+                best_timing,
+                candidate_config,
+                candidate_timing,
+            )
+
+            return True, candidate_timing
+        return False, candidate_timing
+
+    def autotune(
+        self,
+        # pyrefly: ignore [missing-attribute]
+        func: Callable[["triton.Config"], float],
+        # pyrefly: ignore [missing-attribute]
+        baseline_config: "triton.Config",
+        baseline_timing: float | None = None,
+    ) -> "triton.Config":  # pyrefly: ignore  # missing-attribute
+        if baseline_timing is None:
+            baseline_timing = self.call_func(func, baseline_config)
+
+        log.debug("= Do coordinate descent tuning for %s =", self.name)
+        log.debug(
+            "%s: Baseline Config %s, baseline timing %f",
+            self.name,
+            baseline_config,
+            baseline_timing,
+        )
+        improved = True
+        best_config = baseline_config
+        best_timing = baseline_timing
+        tunable_fields = self.tunable_fields
+
+        while improved:
+            improved = False
+
+            for name in tunable_fields:
+                cur_val = get_field(best_config, name)
+                # some kernel don't have R0_BLOCK/YBLOCK/ZBLOCK. So cur_val may be None
+                if cur_val is None:
+                    continue
+
+                # It's possible that candidate_values is empty.
+                # E.g., if XBLOCK is 1 initially and size_hint for x is also 1.
+                # We would not try either larger or smaller XBLOCK in this case.
+                candidate_values = self.get_neighbour_values(name, cur_val)
+
+                for next_val in candidate_values:
+                    candidate_config = copy.deepcopy(best_config)
+                    set_field(candidate_config, name, next_val)
+
+                    if not self.is_valid_config(candidate_config):
+                        continue
+                    cmp_res, candidate_timing = self.compare_config(
+                        func, candidate_config, best_config, best_timing
+                    )
+                    if cmp_res:
+                        improved = True
+                        best_config, best_timing = candidate_config, candidate_timing
+
+            if not improved and self.inductor_meta.get(
+                "coordinate_descent_check_all_directions"
+            ):
+                old_best_timing = best_timing
+                improved, best_config, best_timing = self.check_all_tuning_directions(
+                    func, best_config, best_timing
+                )
+
+                if improved:
+                    msg = red_text(
+                        "%s: Coordinate descend tuning found improvement of %.3fx by looking in all directions."
+                    )
+                    log.debug(
+                        msg,
+                        self.name,
+                        old_best_timing / best_timing,
+                    )
+
+        log.debug(
+            "%s: Improve from %s %f -> %s %f, %.3fx",
+            self.name,
+            baseline_config,
+            baseline_timing,
+            best_config,
+            best_timing,
+            baseline_timing / best_timing,
+        )
+
+        return best_config
+
+    @staticmethod
+    def autotune_single_field(fn, init_val, min_val=None, max_val=None):
+        """
+        fn is a function that takes the field value and returns the benchmarking result
+        init_val is the starting point of autotuning.
+
+        Should work well for parabola like curve. Here is a real example
+        for split-size of mix-order-reduction: https://github.com/pytorch/pytorch/pull/166461
+        """
+        cache = {}
+
+        def _bench(val):
+            if val not in cache:
+                cache[val] = fn(val)
+                # print(f"split size {val} -> {cache[val]:.3f} ms")
+            return cache[val]
+
+        if min_val is None:
+            min_val = 1
+        if max_val is None:
+            max_val = 2**30  # some arbitrary large value
+
+        best_val = init_val
+        improved = True
+        while improved:
+            improved = False
+            candlist = [best_val // 2, best_val * 2]
+            for cand in candlist:
+                cand = max(cand, min_val)
+                cand = min(cand, max_val)
+
+                if _bench(cand) < _bench(best_val):
+                    best_val = cand
+                    improved = True
+
+        return best_val

outputs/audit_venv/lib/python3.11/site-packages/torch/_inductor/runtime/debug_utils.py

@@ -0,0 +1,138 @@
+import functools
+import logging
+import threading
+import weakref
+
+import torch
+from torch.utils._ordered_set import OrderedSet
+
+
+log = logging.getLogger(__name__)
+
+local = threading.local()
+local.memory_tracker = None
+
+
+class BufferMemoryTracker:
+    """
+    Tracks inductor runtime allocations and deallocations to compare against
+    expected behavior.
+    """
+
+    def __init__(self) -> None:
+        self.tensor_tracker: dict[str, torch.storage.UntypedStorage] = (
+            weakref.WeakValueDictionary()  # type: ignore[assignment]
+        )
+        self.died_since_last_step: OrderedSet[str] = OrderedSet()
+        self.added_since_last_step: OrderedSet[str] = OrderedSet()
+        self.error = (
+            torch._inductor.config.test_configs.track_memory_lifecycle == "assert"
+        )
+
+    def set_tensor(self, name: str, tensor: torch.Tensor) -> None:
+        storage = tensor.untyped_storage()
+
+        self.added_since_last_step.add(name)
+        self.tensor_tracker[name] = storage
+
+        def on_tensor_death() -> None:
+            self.died_since_last_step.add(name)
+
+        weakref.finalize(storage, on_tensor_death)
+
+    def advance_step(self) -> None:
+        self.died_since_last_step.clear()
+        self.added_since_last_step.clear()
+
+    def log_or_raise(self, msg: str) -> None:
+        if self.error:
+            raise RuntimeError(msg)
+        else:
+            log.info(msg)
+
+    def check_step_delta(
+        self,
+        expected_allocated: list[str],
+        expected_freed: list[str],
+        is_final_step: bool,
+    ) -> None:
+        """Check only the delta changes since last step"""
+
+        # Check expected deaths - we dont currently distinguish between nodes which die in last step
+        # and are returned as outputs, so skip if final_step.
+        if not is_final_step:
+            missing_deaths = OrderedSet(expected_freed) - self.died_since_last_step
+            if missing_deaths:
+                self.log_or_raise(
+                    f"Expected tensors to die but still alive: {missing_deaths}"
+                )
+
+        # Check for unexpected deaths
+        unexpected_deaths = self.died_since_last_step - OrderedSet(expected_freed)
+        if unexpected_deaths:
+            self.log_or_raise(f"Unexpected tensor deaths: {unexpected_deaths}")
+
+        # Check newly alive tensors - separate messages like deaths
+        actual_allocated = self.added_since_last_step
+        expected_allocated_set = OrderedSet(expected_allocated)
+
+        extra_alive = actual_allocated - expected_allocated_set
+        if extra_alive:
+            self.log_or_raise(f"Unexpected allocated tensors: {extra_alive}")
+
+        missing_alive = expected_allocated_set - actual_allocated
+        if missing_alive:
+            self.log_or_raise(
+                f"Expected allocated tensors but missing: {missing_alive}"
+            )
+
+        # Reset for next step
+        self.advance_step()
+
+        if is_final_step:
+            local.memory_tracker = None
+
+
+def get_mem_tracker() -> BufferMemoryTracker:
+    if local.memory_tracker is None:
+        local.memory_tracker = BufferMemoryTracker()
+    return local.memory_tracker
+
+
+def track_tensor(tensor: torch.Tensor, name: str) -> None:
+    get_mem_tracker().set_tensor(name, tensor)
+
+
+def tracked_empty_strided(
+    size: list[int],
+    stride: list[int],
+    *,
+    dtype: torch.dtype,
+    device: torch.device,
+    name: str,
+) -> torch.Tensor:
+    o = torch.empty_strided(size, stride, dtype=dtype, device=device)
+    track_tensor(o, name)
+    return o
+
+
+def check_memory_step(
+    allocated: list[str], freed: list[str], is_final_step: bool = False
+) -> None:
+    tracker = get_mem_tracker()
+    tracker.check_step_delta(allocated, freed, is_final_step)
+
+
+@functools.lru_cache(None)
+def register_check_mem_op() -> None:
+    lib = torch.library.Library("_inductor_debug", "FRAGMENT")  # noqa: TOR901
+    lib.define(
+        "check_memory_step(str[] allocated, str[] freed, bool is_final_step) -> ()"
+    )
+    lib.impl("check_memory_step", check_memory_step, "BackendSelect")
+    from torch._higher_order_ops.effects import _EffectType, _register_effectful_op
+
+    _register_effectful_op(
+        torch.ops._inductor_debug.check_memory_step.default,
+        _EffectType.ORDERED,
+    )

outputs/audit_venv/lib/python3.11/site-packages/torch/_inductor/runtime/halide_helpers.py

@@ -0,0 +1,155 @@
+# mypy: allow-untyped-defs
+import math
+
+
+try:
+    import halide as hl  # type: ignore[import-untyped, import-not-found]
+except ImportError:
+    hl = None
+
+PHILOX_N_ROUNDS_DEFAULT = 10  # Default number of rounds for philox
+
+if hl is not None:
+    PHILOX_KEY_A_U32 = hl.u32(0x9E3779B9)
+    PHILOX_KEY_B_U32 = hl.u32(0xBB67AE85)
+    PHILOX_ROUND_A_U32 = hl.u32(0xD2511F53)
+    PHILOX_ROUND_B_U32 = hl.u32(0xCD9E8D57)
+else:
+    PHILOX_KEY_A_U32 = None
+    PHILOX_KEY_B_U32 = None
+    PHILOX_ROUND_A_U32 = None
+    PHILOX_ROUND_B_U32 = None
+
+
+def _pair_uniform_to_normal(u1, u2):
+    """Box-Muller transform"""
+    u1 = hl.max(hl.f32(1.0e-7), u1)
+    th = hl.f32(math.tau) * u2
+    r = hl.sqrt(hl.f32(-2.0) * hl.log(u1))
+    return r * hl.cos(th), r * hl.sin(th)
+
+
+def _uint_to_uniform_float(x):
+    """
+    Numerically stable function to convert a random uint into a random float uniformly sampled in [0, 1).
+    """
+
+    # TODO:
+    # conditions can be simplified
+    # scale is ((2**23 - 1) / 2**23) * 2**(N_BITS - 1)
+    # https://github.com/triton-lang/triton/blob/e4a0d93ff1a367c7d4eeebbcd7079ed267e6b06f/python/triton/language/random.py#L116-L132.
+    assert x.type() == hl.UInt(32) or x.type() == hl.Int(32)
+    x = hl.cast(hl.Int(32), x)
+    scale = hl.f64(4.6566127342e-10)
+    x = hl.select(x < 0, -x - 1, x)
+    return x * scale
+
+
+def philox_impl(c0, c1, c2, c3, k0, k1, n_rounds):
+    def umulhi(a, b):
+        a = hl.cast(hl.UInt(64), a)
+        b = hl.cast(hl.UInt(64), b)
+        return hl.cast(hl.UInt(32), ((a * b) >> 32) & hl.u64(0xFFFFFFFF))
+
+    for _ in range(n_rounds):
+        _c0, _c2 = c0, c2
+
+        c0 = umulhi(PHILOX_ROUND_B_U32, _c2) ^ c1 ^ k0
+        c2 = umulhi(PHILOX_ROUND_A_U32, _c0) ^ c3 ^ k1
+        c1 = PHILOX_ROUND_B_U32 * _c2
+        c3 = PHILOX_ROUND_A_U32 * _c0
+        # raise key
+        k0 = k0 + PHILOX_KEY_A_U32
+        k1 = k1 + PHILOX_KEY_B_U32
+
+    return c0, c1, c2, c3
+
+
+def halide_philox(seed, c0, c1, c2, c3, n_rounds):
+    seed = hl.cast(hl.UInt(64), seed)
+
+    assert c0.type().bits() == 32
+
+    seed_hi = hl.cast(hl.UInt(32), (seed >> 32) & hl.u64(0xFFFFFFFF))
+    seed_lo = hl.cast(hl.UInt(32), seed & hl.u64(0xFFFFFFFF))
+
+    return philox_impl(c0, c1, c2, c3, seed_lo, seed_hi, n_rounds)
+
+
+def randint4x(seed, offset, n_rounds):
+    offset = hl.cast(hl.UInt(32), offset)
+    _0 = hl.u32(0)
+    return halide_philox(seed, offset, _0, _0, _0, n_rounds)
+
+
+def rand4x(seed, offset, n_rounds=PHILOX_N_ROUNDS_DEFAULT):
+    i1, i2, i3, i4 = randint4x(seed, offset, n_rounds)
+    u1 = _uint_to_uniform_float(i1)
+    u2 = _uint_to_uniform_float(i2)
+    u3 = _uint_to_uniform_float(i3)
+    u4 = _uint_to_uniform_float(i4)
+    return u1, u2, u3, u4
+
+
+def randint(seed, offset, n_rounds=PHILOX_N_ROUNDS_DEFAULT):
+    ret, _, _, _ = randint4x(seed, offset, n_rounds)
+    return ret
+
+
+def rand(seed, offset, n_rounds=PHILOX_N_ROUNDS_DEFAULT):
+    source = randint(seed, offset, n_rounds)
+    return _uint_to_uniform_float(source)
+
+
+def rand_eager_kernel(seed, offset_blocks, tid, VEC, n_rounds=PHILOX_N_ROUNDS_DEFAULT):
+    inv = hl.cast(hl.Float(32), 1.0 / 4294967296.0)  # 2^-32
+    half = hl.cast(hl.Float(32), 0.5) * inv
+
+    tid_u64 = hl.cast(hl.UInt(64), tid)
+    VEC_u64 = hl.cast(hl.UInt(64), VEC)
+    subseq = tid_u64 // VEC_u64
+    which4 = (tid_u64 % VEC_u64) // hl.cast(hl.UInt(64), 4)
+    lane = tid_u64 % hl.cast(hl.UInt(64), 4)
+
+    offblk = hl.cast(hl.UInt(64), offset_blocks) + which4
+
+    c0 = hl.cast(hl.UInt(32), offblk & hl.cast(hl.UInt(64), 0xFFFFFFFF))
+    c1 = hl.cast(
+        hl.UInt(32),
+        (offblk >> hl.cast(hl.UInt(64), 32)) & hl.cast(hl.UInt(64), 0xFFFFFFFF),
+    )
+    c2 = hl.cast(hl.UInt(32), subseq & hl.cast(hl.UInt(64), 0xFFFFFFFF))
+    c3 = hl.cast(
+        hl.UInt(32),
+        (subseq >> hl.cast(hl.UInt(64), 32)) & hl.cast(hl.UInt(64), 0xFFFFFFFF),
+    )
+
+    u0, u1, u2, u3 = halide_philox(seed, c0, c1, c2, c3, n_rounds)
+
+    v01 = hl.select(lane == hl.cast(hl.UInt(64), 0), u0, u1)
+    v23 = hl.select(lane == hl.cast(hl.UInt(64), 2), u2, u3)
+    rand_int = hl.select(
+        (lane == hl.cast(hl.UInt(64), 0)) | (lane == hl.cast(hl.UInt(64), 1)), v01, v23
+    )
+
+    return hl.cast(hl.Float(32), 1.0) - (hl.cast(hl.Float(32), rand_int) * inv + half)
+
+
+def randn(seed, offset):
+    i1, i2, _, _ = randint4x(seed, offset, PHILOX_N_ROUNDS_DEFAULT)
+    u1 = _uint_to_uniform_float(i1)
+    u2 = _uint_to_uniform_float(i2)
+    n1, _ = _pair_uniform_to_normal(u1, u2)
+    return n1
+
+
+def randint64(seed, offset, low, high):
+    r0, r1, _r2, _r3 = randint4x(seed, offset, PHILOX_N_ROUNDS_DEFAULT)
+    r0 = hl.cast(hl.UInt(64), r0)
+    r1 = hl.cast(hl.UInt(64), r1)
+
+    result = r0 | (r1 << 32)
+    size = high - low
+    result = result % hl.cast(hl.UInt(64), size)
+    result = hl.cast(hl.Int(64), result) + low
+    return result

outputs/audit_venv/lib/python3.11/site-packages/torch/_inductor/runtime/hints.py

@@ -0,0 +1,236 @@
+# mypy: allow-untyped-defs
+from __future__ import annotations
+
+import collections
+import functools
+import typing
+from enum import auto, Enum
+
+import torch
+from torch.utils._triton import has_triton_package
+
+
+# The following maximums only apply to runtime autotuning, when using FixedTritonConfig one may see larger values
+# NOTE: if these fail asserts submit a PR to increase them
+TRITON_MAX_BLOCK = {
+    "X": 8192 if torch.version.hip else 4096,
+    "Y": 1024,
+    "Z": 1024,
+    "R0_": 4096 * 16,  # * 16 is multi-kernel only
+    "R1_": 2048 * 16,  # * 16 is multi-kernel only
+}
+TRITON_MAX_RSPLIT = 64
+
+
+class ReductionHint(Enum):
+    INNER = 0
+    OUTER = 1
+    OUTER_TINY = 2
+    DEFAULT = 3
+
+
+class TileHint(Enum):
+    SQUARE = 0
+    DEFAULT = 1
+
+
+# Define `AttrsDescriptorWrapper` function with clear conditional handling
+if has_triton_package():
+    import triton
+    import triton.backends.compiler
+    import triton.compiler.compiler
+
+    if hasattr(triton.backends.compiler, "AttrsDescriptor"):
+        # Triton 3.2.0 - the second implementation
+        from triton.backends.compiler import AttrsDescriptor
+
+        def AttrsDescriptorWrapper(
+            divisible_by_16=None,
+            equal_to_1=None,
+            pointer_range_32=None,
+        ):
+            # Prepare the arguments for AttrsDescriptor
+            kwargs = {
+                "tt.divisibility": divisible_by_16,
+                "tt.equal_to": equal_to_1,
+            }
+
+            # Instantiate AttrsDescriptor with the prepared arguments
+            res = AttrsDescriptor.from_dict(
+                {"arg_properties": kwargs, "cls": AttrsDescriptor.__name__}
+            )
+            assert res.property_values["tt.divisibility"] == 16
+            assert res.property_values["tt.equal_to"] == 1
+            return res
+
+    elif hasattr(triton.compiler.compiler, "AttrsDescriptor"):
+        # Triton 3.0.0 - the original implementation
+        from triton.compiler.compiler import AttrsDescriptor
+
+        def AttrsDescriptorWrapper(
+            divisible_by_16=None,
+            equal_to_1=None,
+            pointer_range_32=None,
+        ):
+            # Prepare the arguments for AttrsDescriptor
+            kwargs = {
+                "divisible_by_16": divisible_by_16,
+                "equal_to_1": equal_to_1,
+            }
+
+            # Instantiate AttrsDescriptor with the prepared arguments
+            return AttrsDescriptor(**kwargs)
+
+    else:
+        # Triton in 2025:
+        # note: there's also a range of triton commits not currently supported
+        # from ~Dec 9, 2024 to Jan 1 2025, in which AttrsDescriptors are still
+        # used, but the contents are different.
+
+        def AttrsDescriptorWrapper(
+            divisible_by_16=None,
+            equal_to_1=None,
+            pointer_range_32=None,
+        ):
+            # pyrefly: ignore [not-iterable]
+            # Build attr dict merging divisibility and pointer_range per arg index,
+            # since a single arg can carry both attributes.
+            result = {(x,): [["tt.divisibility", 16]] for x in (divisible_by_16 or ())}
+            for x in pointer_range_32 or ():
+                key = (x,)
+                if key in result:
+                    result[key].append(["tt.pointer_range", 32])
+                else:
+                    result[key] = [["tt.pointer_range", 32]]
+            return result
+
+else:
+    # Define a namedtuple as a fallback when AttrsDescriptor is not available
+    AttrsDescriptorWrapper = collections.namedtuple(  # type: ignore[no-redef, name-match]
+        # pyrefly: ignore [invalid-argument]
+        "AttrsDescriptor",
+        ["divisible_by_16", "equal_to_1", "pointer_range_32"],
+        defaults=[(), (), ()],
+    )
+
+
+_NUM_THREADS_PER_WARP = 32
+
+
+class HeuristicType(Enum):
+    PERSISTENT_REDUCTION = auto()
+    POINTWISE = auto()
+    REDUCTION = auto()
+    SPLIT_SCAN = auto()
+    TEMPLATE = auto()
+    USER_AUTOTUNE = auto()
+    FIXED = auto()
+
+
+class AutotuneHint(Enum):
+    ONE_ELEMENT_PER_THREAD = 0
+
+    # Triton codegen tries to codegen set of AutotuneHints.
+    # Enum.__repr__ looks like "<AutotuneHint.ELEMENTS_PER_WARP_32: 0>""
+    # which isn't valid python.
+    # Enum.__str__ will just return "AutotuneHint.ELEMENTS_PER_WARP_32".
+    __repr__ = Enum.__str__
+
+
+class DeviceProperties(typing.NamedTuple):
+    """Copy device properties into a data structure not requiring torch to be imported"""
+
+    type: str  # type: ignore[assignment]
+    index: int  # type: ignore[assignment]
+    multi_processor_count: int
+    cc: int
+    major: int | None = None
+    regs_per_multiprocessor: int | None = None
+    max_threads_per_multi_processor: int | None = None
+    max_threads_per_block: int | None = None
+    warp_size: int | None = None
+
+    @classmethod
+    @functools.cache
+    def create(cls, device) -> DeviceProperties:
+        import torch
+        from torch._dynamo.device_interface import get_interface_for_device
+
+        device_type = device.type
+
+        if torch.version.hip and device_type == "cuda":
+            device_type = "hip"
+
+        device_interface = get_interface_for_device(device)
+        props = device_interface.get_device_properties(device)
+        try:
+            multi_processor_count = props.multi_processor_count
+        except AttributeError:
+            if device_type == "xpu":
+                multi_processor_count = props.gpu_subslice_count
+            elif device_type == "mtia":
+                multi_processor_count = 64
+            else:
+                raise
+        return cls(
+            type=device_type,
+            index=device.index,
+            multi_processor_count=multi_processor_count,
+            cc=device_interface.get_compute_capability(device),
+            major=getattr(props, "major", None),
+            regs_per_multiprocessor=getattr(props, "regs_per_multiprocessor", None),
+            max_threads_per_multi_processor=getattr(
+                props, "max_threads_per_multi_processor", None
+            ),
+            max_threads_per_block=getattr(props, "max_threads_per_block", 1024),
+            warp_size=getattr(props, "warp_size", 32 if device_type != "cpu" else None),
+        )
+
+
+class HalideInputSpec(typing.NamedTuple):
+    ctype: str
+    name: str
+    shape: list[str] | None = None
+    stride: list[str] | None = None
+    offset: str | None = None
+    alias_of: str | None = None
+
+    def bindings_type(self) -> str:
+        if self.ctype in ("at::Half*", "at::BFloat16*"):
+            return "uint16_t*"  # half not defined
+        return self.ctype
+
+    def halide_type(self) -> str:
+        if self.ctype == "at::Half*":
+            return "halide_type_t(halide_type_float, 16)"  # half not defined
+        if self.ctype == "at::BFloat16*":
+            return "halide_type_t(halide_type_bfloat, 16)"  # half not defined
+        return f"halide_type_of<{self.ctype.replace('*', '')}>()"
+
+    def is_scalar(self) -> bool:
+        return self.shape is None
+
+    def is_buffer(self) -> bool:
+        return self.shape is not None
+
+
+class HalideMeta(typing.NamedTuple):
+    argtypes: list[HalideInputSpec]
+    target: str
+    scheduler: str | None = None
+    scheduler_flags: dict[str, int | str] | None = None
+    cuda_device: int | None = None
+
+    def args(self) -> list[str]:
+        """Command line args to pass to halide generator"""
+        args = [f"target={self.target}"]
+        if self.scheduler:
+            args.append(f"autoscheduler={self.scheduler}")
+        if self.scheduler_flags:
+            assert self.scheduler
+            for k, v in self.scheduler_flags.items():
+                args.append(f"autoscheduler.{k}={v}")
+        return args
+
+    def is_cuda(self) -> bool:
+        return self.cuda_device is not None

outputs/audit_venv/lib/python3.11/site-packages/torch/_inductor/runtime/proton_utils.py

@@ -0,0 +1,325 @@
+"""
+Utilities for post-processing Triton Proton profiling trace files.
+
+This module provides functions to transform proton-generated Chrome trace files
+for better visualization in Perfetto/Chrome tracing tools.
+"""
+
+import gzip
+import json
+import os
+import re
+from typing import Any
+
+from torch.utils._ordered_set import OrderedSet
+
+
+def _write_trace(path: str, data: dict[str, Any]) -> None:
+    """Write trace data to a gzipped JSON file."""
+    with gzip.open(path, "wt", encoding="utf-8") as f:
+        json.dump(data, f)
+
+
+def _read_trace(path: str) -> dict[str, Any]:
+    """Read trace data from a JSON file (gzipped or not)."""
+    if path.endswith(".gz"):
+        with gzip.open(path, "rt", encoding="utf-8") as f:
+            return json.load(f)
+    else:
+        with open(path) as f:
+            return json.load(f)
+
+
+def _get_base_name(trace_path: str) -> str:
+    """Get the base name of a trace file, stripping known extensions."""
+    base_name = os.path.basename(trace_path)
+    for ext in [".trace.json.gz", ".chrome_trace", ".json.gz", ".json"]:
+        if base_name.endswith(ext):
+            return base_name[: -len(ext)]
+    return os.path.splitext(base_name)[0]
+
+
+def _group_events_by_sm(events: list[dict[str, Any]]) -> list[dict[str, Any]]:
+    """
+    Group events so all CTAs on the same SM share one track.
+
+    Transforms:
+    - pid: "Core0 CTA6" -> "Core0"
+    - tid: "warp0" -> "CTA6 warp0"
+    """
+    core_cta_pattern = re.compile(r"^(.*?)\s*(Core\d+)\s+(CTA\d+)$")
+
+    for event in events:
+        pid = event.get("pid", "")
+        tid = event.get("tid", "")
+
+        match = core_cta_pattern.match(str(pid))
+        if match:
+            prefix = match.group(1).strip()
+            core = match.group(2)
+            cta = match.group(3)
+            event["pid"] = f"{prefix} {core}" if prefix else core
+            event["tid"] = f"{cta} {tid}" if tid else cta
+
+    return events
+
+
+def _group_events_per_cta_occupancy(
+    events: list[dict[str, Any]],
+) -> list[dict[str, Any]]:
+    """
+    Process warp events into CTA-level events with slot assignment.
+
+    Each CTA gets: min warp start, max warp end.
+    CTAs are assigned to slots per SM such that non-overlapping CTAs share slots.
+    """
+    core_cta_pattern = re.compile(r"^(.*?)\s*(Core\d+)\s+(CTA\d+)$")
+
+    # Group events by (prefix, SM, CTA)
+    cta_events: dict[tuple[str, str, str], list[dict[str, Any]]] = {}
+
+    for event in events:
+        pid = event.get("pid", "")
+        match = core_cta_pattern.match(str(pid))
+        if match:
+            prefix = match.group(1).strip()
+            core = match.group(2)
+            cta = match.group(3)
+            key = (prefix, core, cta)
+            if key not in cta_events:
+                cta_events[key] = []
+            cta_events[key].append(event)
+
+    # For each CTA, compute min start and max end
+    cta_intervals: dict[tuple[str, str, str], tuple[float, float]] = {}
+    for key, evts in cta_events.items():
+        min_start = float("inf")
+        max_end = float("-inf")
+        for evt in evts:
+            ts = evt.get("ts", 0)
+            dur = evt.get("dur", 0)
+            min_start = min(min_start, ts)
+            max_end = max(max_end, ts + dur)
+        if min_start != float("inf"):
+            cta_intervals[key] = (min_start, max_end)
+
+    # Group CTAs by (prefix, SM) and assign to slots
+    sm_ctas: dict[tuple[str, str], list[tuple[str, float, float]]] = {}
+    for (prefix, core, cta), (start, end) in cta_intervals.items():
+        sm_key = (prefix, core)
+        if sm_key not in sm_ctas:
+            sm_ctas[sm_key] = []
+        sm_ctas[sm_key].append((cta, start, end))
+
+    # Assign CTAs to slots using interval scheduling (greedy)
+    cta_slot_assignments: dict[tuple[str, str, str], int] = {}
+    for sm_key, ctas in sm_ctas.items():
+        prefix, core = sm_key
+        sorted_ctas = sorted(ctas, key=lambda x: x[1])
+        slots: list[float] = []
+        for cta, start, end in sorted_ctas:
+            assigned_slot = None
+            for i, slot_end in enumerate(slots):
+                if start >= slot_end:
+                    assigned_slot = i
+                    slots[i] = end
+                    break
+            if assigned_slot is None:
+                assigned_slot = len(slots)
+                slots.append(end)
+            cta_slot_assignments[(prefix, core, cta)] = assigned_slot
+
+    # Build numeric ID mappings for pids and tids
+    # Chrome trace format requires numeric pid/tid values
+    pid_names: dict[str, int] = {}
+    tid_names: dict[tuple[int, str], int] = {}  # (pid, tid_name) -> tid_num
+
+    for key in cta_intervals:
+        prefix, core, cta = key
+        pid_name = f"{prefix} {core}" if prefix else core
+        if pid_name not in pid_names:
+            pid_names[pid_name] = len(pid_names)
+
+    new_events: list[dict[str, Any]] = []
+
+    # Add process name metadata events
+    for pid_name, pid_num in pid_names.items():
+        new_events.append(
+            {
+                "name": "process_name",
+                "ph": "M",
+                "pid": pid_num,
+                "args": {"name": pid_name},
+            }
+        )
+
+    # Create CTA events with numeric IDs and collect tid mappings
+    for key, (start, end) in cta_intervals.items():
+        prefix, core, cta = key
+        slot = cta_slot_assignments[key]
+        pid_name = f"{prefix} {core}" if prefix else core
+        tid_name = f"slot{slot}"
+        pid_num = pid_names[pid_name]
+
+        if (pid_num, tid_name) not in tid_names:
+            tid_num = len(tid_names)
+            tid_names[(pid_num, tid_name)] = tid_num
+            # Add thread name metadata event
+            new_events.append(
+                {
+                    "name": "thread_name",
+                    "ph": "M",
+                    "pid": pid_num,
+                    "tid": tid_num,
+                    "args": {"name": tid_name},
+                }
+            )
+        else:
+            tid_num = tid_names[(pid_num, tid_name)]
+
+        new_events.append(
+            {
+                "name": cta,
+                "cat": "cta",
+                "ph": "X",
+                "ts": start,
+                "dur": end - start,
+                "pid": pid_num,
+                "tid": tid_num,
+            }
+        )
+
+    return new_events
+
+
+def _apply_grouping(
+    events: list[dict[str, Any]], group_by_sm: bool, per_cta_occupancy: bool
+) -> list[dict[str, Any]]:
+    """Apply grouping transformation to events."""
+    if per_cta_occupancy:
+        return _group_events_per_cta_occupancy(events)
+    elif group_by_sm:
+        return _group_events_by_sm(events)
+    return events
+
+
+def _split_events_by_invocation(
+    events: list[dict[str, Any]],
+    gap_threshold_ns: float = 1000.0,
+) -> list[list[dict[str, Any]]]:
+    """Split events into separate invocations based on time gaps."""
+    if not events:
+        return []
+
+    events_sorted = sorted(events, key=lambda e: e.get("ts", 0))
+    invocations: list[list[dict[str, Any]]] = [[]]
+    prev_end = events_sorted[0].get("ts", 0)
+
+    for event in events_sorted:
+        ts = event.get("ts", 0)
+        dur = event.get("dur", 0)
+        if ts - prev_end > gap_threshold_ns and invocations[-1]:
+            invocations.append([])
+        invocations[-1].append(event)
+        prev_end = max(prev_end, ts + dur)
+
+    return [inv for inv in invocations if inv]
+
+
+def _normalize_timestamps(
+    events: list[dict[str, Any]], scale_factor: float = 1.0
+) -> list[dict[str, Any]]:
+    """Normalize timestamps to start at 0 and optionally scale."""
+    if not events:
+        return events
+
+    # Only consider duration events (ph=X) for min timestamp calculation
+    duration_events = [e for e in events if e.get("ph") == "X"]
+    if not duration_events:
+        return events
+
+    min_ts = min(e.get("ts", 0) for e in duration_events)
+    for event in events:
+        if event.get("ph") == "X":
+            event["ts"] = (event.get("ts", 0) - min_ts) * scale_factor
+            if "dur" in event:
+                event["dur"] = event["dur"] * scale_factor
+
+    return events
+
+
+def process_proton_trace(
+    trace_path: str,
+    output_dir: str | None = None,
+    group_by_sm: bool = True,
+    split_invocations: bool = True,
+    scale_factor: float = 1.0,
+    gap_threshold_ns: float = 1000.0,
+    per_cta_occupancy: bool = True,
+) -> list[str]:
+    """
+    Process a proton trace file with various transformations.
+
+    Always produces 1 + N files: one complete trace and N per-invocation traces.
+    Grouping options (group_by_sm, per_cta_occupancy) apply uniformly to all outputs.
+
+    Args:
+        trace_path: Path to the input Chrome trace file
+        output_dir: Directory to write output files. If None, uses same directory as input.
+        group_by_sm: If True, group CTAs by SM into single tracks.
+        split_invocations: If True, also produce per-invocation trace files.
+        scale_factor: Factor to scale durations by (helps visibility in Perfetto).
+        gap_threshold_ns: Time gap (in nanoseconds) that indicates a new invocation.
+        per_cta_occupancy: If True, process warp tracks into CTA tracks and assign
+            CTAs to slots per SM such that CTAs do not overlap.
+
+    Returns:
+        List of paths to the output files.
+    """
+    if output_dir is None:
+        output_dir = os.path.dirname(trace_path) or "."
+
+    os.makedirs(output_dir, exist_ok=True)
+    base_name = _get_base_name(trace_path)
+
+    data = _read_trace(trace_path)
+    events = data.get("traceEvents", [])
+
+    # Split into invocations first (before grouping, to avoid merging across invocations)
+    invocations = _split_events_by_invocation(events, gap_threshold_ns)
+
+    # Apply grouping transformation to each invocation
+    invocations = [
+        _apply_grouping(inv_events, group_by_sm, per_cta_occupancy)
+        for inv_events in invocations
+    ]
+
+    output_files = []
+
+    # Write complete trace (dedupe metadata events)
+    complete_events = []
+    seen_metadata: OrderedSet[tuple[str | None, int | None, int | None]] = OrderedSet()
+    for inv_events in invocations:
+        for event in inv_events:
+            if event.get("ph") == "M":
+                # Dedupe metadata events by (name, pid, tid)
+                key = (event.get("name"), event.get("pid"), event.get("tid"))
+                if key in seen_metadata:
+                    continue
+                seen_metadata.add(key)
+            complete_events.append(event)
+    complete_path = os.path.join(output_dir, f"{base_name}.trace.json.gz")
+    _write_trace(complete_path, {"traceEvents": complete_events})
+    output_files.append(complete_path)
+
+    # Write per-invocation traces
+    if split_invocations:
+        for i, inv_events in enumerate(invocations):
+            inv_events = _normalize_timestamps(list(inv_events), scale_factor)
+            inv_path = os.path.join(
+                output_dir, f"{base_name}_invocation_{i}.trace.json.gz"
+            )
+            _write_trace(inv_path, {"traceEvents": inv_events})
+            output_files.append(inv_path)
+
+    return output_files

outputs/audit_venv/lib/python3.11/site-packages/torch/_inductor/runtime/runtime_utils.py

@@ -0,0 +1,829 @@
+from __future__ import annotations
+
+import functools
+import math
+import operator
+from typing import Any, TYPE_CHECKING
+
+import sympy
+
+import torch
+
+# NOTE: other files rely on the imports below
+from torch._dynamo import callback as compilation_callback  # noqa: F401
+from torch._inductor.runtime.cache_dir_utils import (  # noqa: F401
+    cache_dir,
+    default_cache_dir,
+    triton_cache_dir,
+)
+
+
+if TYPE_CHECKING:
+    from collections.abc import Hashable
+
+    from .triton_compat import Config
+
+
+def conditional_product(*args: int) -> int:
+    return functools.reduce(operator.mul, [x for x in args if x])
+
+
+def ceildiv(number: int, denom: int) -> int:
+    return -(number // -denom)
+
+
+def is_power_of_2(n: int) -> bool:
+    """Returns whether n = 2 ** m for some integer m."""
+    return n > 0 and n & n - 1 == 0
+
+
+def next_power_of_2(n: int) -> int:
+    """Return the smallest power of 2 greater than or equal to n"""
+    if isinstance(n, sympy.Integer):
+        n = int(n)
+    if n <= 0:
+        return 1
+    return 1 << (n - 1).bit_length()
+
+
+def last_power_of_2(n: int) -> int:
+    """Return the largest power of 2 less than or equal to n"""
+    next_pow2 = next_power_of_2(n)
+    return next_pow2 // 2 if next_pow2 > n else next_pow2
+
+
+def get_num_bytes(*args: torch.Tensor, num_in_out_args: int = 0) -> int:
+    """
+    Return the total number of bytes the arguments of tensor type takes.
+
+    For in/out args, tensor sizes are counted twice: once for reading and
+    once for writing.
+
+    The first num_in_out_args arguments are in out tensors.
+    """
+    return sum(
+        arg.numel() * arg.element_size() * (1 + int(i < num_in_out_args))
+        for i, arg in enumerate(args)
+        if isinstance(arg, torch.Tensor)
+    )
+
+
+def triton_config_to_hashable(cfg: Config) -> Hashable:
+    """
+    Convert triton config to a tuple that can uniquely identify it. We can use
+    the return value as a dictionary key.
+    """
+    # pyrefly: ignore [missing-attribute]
+    items = sorted(cfg.kwargs.items())
+    # pyrefly: ignore [missing-attribute]
+    items.append(("num_warps", cfg.num_warps))
+    # pyrefly: ignore [missing-attribute]
+    items.append(("num_stages", cfg.num_stages))
+    return tuple(items)
+
+
+def validate_triton_config(cfg: Config) -> None:
+    # [Note: Triton pre_hook in inductor]
+    # pre-hook is a lambda function, which we don't attempt to serialize.
+    # right now, if a pre-hook is attached to the config, it will not be saved;
+    # and then it won't be used when the config is loaded from cache.
+    # So we assert - if we do get a pre_hook, it might get ignored after caching.
+    assert getattr(cfg, "pre_hook", None) is None, (
+        "triton configs with pre_hooks not supported"
+    )
+
+
+def create_bandwidth_info_str(
+    ms: float,
+    num_gb: float,
+    gb_per_s: float,
+    prefix: str = "",
+    suffix: str = "",
+    color: bool = True,
+) -> str:
+    info_str = f"{prefix}{ms:.3f}ms    \t{num_gb:.3f} GB \t {gb_per_s:7.2f}GB/s{suffix}"
+    slow = ms > 0.012 and gb_per_s < 650
+    return red_text(info_str) if color and slow else info_str
+
+
+def get_max_y_grid() -> int:
+    return 65535
+
+
+try:
+    import colorama
+
+    HAS_COLORAMA = True
+except ModuleNotFoundError:
+    HAS_COLORAMA = False
+    colorama = None  # type: ignore[assignment]
+
+
+if HAS_COLORAMA:
+
+    def _color_text(msg: str, color: str) -> str:
+        # pyrefly: ignore [missing-attribute]
+        return getattr(colorama.Fore, color.upper()) + msg + colorama.Fore.RESET
+
+else:
+
+    def _color_text(msg: str, color: str) -> str:
+        return msg
+
+
+def green_text(msg: str) -> str:
+    return _color_text(msg, "green")
+
+
+def yellow_text(msg: str) -> str:
+    return _color_text(msg, "yellow")
+
+
+def red_text(msg: str) -> str:
+    return _color_text(msg, "red")
+
+
+def blue_text(msg: str) -> str:
+    return _color_text(msg, "blue")
+
+
+def get_first_attr(obj: Any, *attrs: str) -> Any:
+    """
+    Return the first available attribute or throw an exception if none is present.
+    """
+    for attr in attrs:
+        if hasattr(obj, attr):
+            return getattr(obj, attr)
+
+    raise AssertionError(f"{obj} does not has any of the attributes: {attrs}")
+
+
+dynamo_timed = torch._dynamo.utils.dynamo_timed  # type: ignore[has-type]
+
+
+def triton_hash_to_path_key(key: str) -> str:
+    # In early versions of Triton, the hash is directly used in the path name.
+    # Later, the hash is converted to base64 before being used in the path name.
+    # Later, the base64 conversion was replaced to the base32
+    #
+    # This code tries to import _base64 and falls back to _base32 if _base64 is unavailable.
+    #
+    # To handle this, try to import the to-base64-conversion function.
+    # If it exists, use it; otherwise, try using _base32; if both are unavailable, use the hash directly.
+    try:
+        from triton.runtime.cache import _base64
+
+        return _base64(key)
+    except Exception:
+        try:
+            from triton.runtime.cache import _base32
+
+            return _base32(key)
+        except Exception:
+            return key
+
+
+def compile_mps_shader(source: str) -> Any:
+    """
+    Compiles shader source but raise more actionable error message when needed
+    """
+    try:
+        return torch.mps.compile_shader(source)
+    except SyntaxError as err:
+        raise SyntaxError(f"failed to compile {source} with {err.msg}") from err
+
+
+def compile_mps_shaders(
+    kernels: list[tuple[str, str, list[str]]],
+) -> dict[str, Any]:
+    """Compile a batch of Metal kernels into one library.
+
+    Args:
+        kernels: list of (kernel_name, metal_source, headers) tuples.
+            headers are bare names resolved as <c10/metal/{name}.h>.
+
+    Returns:
+        dict mapping each kernel_name to its compiled function handle.
+    """
+    from torch.utils._ordered_set import OrderedSet
+
+    all_headers = sorted(OrderedSet(h for _, _, hs in kernels for h in hs))
+    header_src = "\n".join(f"#include <c10/metal/{h}.h>" for h in all_headers)
+    body_src = "\n".join(src for _, src, _ in kernels)
+    lib = compile_mps_shader(header_src + "\n" + body_src)
+    return {name: getattr(lib, name) for name, _, _ in kernels}
+
+
+def torch_dtype_to_jax_runtime(dtype: torch.dtype) -> Any:
+    """
+    Map PyTorch dtype to actual JAX dtype object at runtime.
+
+    This helper is used in generated Pallas kernels at runtime to convert
+    PyTorch dtypes to JAX dtype objects (not string representations).
+
+    Args:
+        dtype: PyTorch dtype to convert
+
+    Returns:
+        JAX dtype object (e.g., jnp.float32 object itself)
+    """
+    import jax.numpy as jnp  # pyrefly: ignore [import-error, missing-import]
+
+    dtype_map = {
+        torch.float32: jnp.float32,
+        torch.float64: jnp.float64,
+        torch.float16: jnp.float16,
+        torch.bfloat16: jnp.bfloat16,
+        torch.int32: jnp.int32,
+        torch.int64: jnp.int64,
+        torch.int16: jnp.int16,
+        torch.int8: jnp.int8,
+        torch.uint8: jnp.uint8,
+        torch.bool: jnp.bool_,
+        torch.complex64: jnp.complex64,
+        torch.complex128: jnp.complex128,
+    }
+    if dtype not in dtype_map:
+        raise ValueError(f"Unsupported dtype for JAX conversion: {dtype}")
+    return dtype_map[dtype]
+
+
+def torch_dtype_to_jax(dtype: torch.dtype) -> str:
+    """
+    Map PyTorch dtype to JAX dtype expression string.
+
+    This helper is used at compile time in codegen to generate
+    JAX dtype expressions for Pallas kernels.
+
+    Args:
+        dtype: PyTorch dtype to convert
+
+    Returns:
+        JAX dtype expression as string (e.g., "jnp.float32")
+    """
+    jax_dtype = torch_dtype_to_jax_runtime(dtype)
+    dtype_name = jax_dtype.__name__
+    if dtype_name == "bool":
+        dtype_name = "bool_"
+    return f"jnp.{dtype_name}"
+
+
+def pallas_gpu_pad_inputs(inputs: list[Any], alignment: int = 128) -> list[Any]:
+    """Flatten and pad each input JAX array to a multiple of alignment."""
+    import jax.numpy as jnp  # pyrefly: ignore [import-error, missing-import]
+
+    padded = []
+    for inp in inputs:
+        flat = inp.flatten()
+        orig_size = flat.size
+        aligned_size = ((orig_size + alignment - 1) // alignment) * alignment
+        if orig_size != aligned_size:
+            padded.append(jnp.pad(flat, (0, aligned_size - orig_size)))
+        else:
+            padded.append(flat)
+    return padded
+
+
+def pallas_gpu_align_output_specs(
+    out_shapes: tuple[Any, ...],
+    out_dtypes: tuple[Any, ...],
+    alignment: int = 128,
+) -> tuple[tuple[Any, ...], list[bool]]:
+    """Build aligned output ShapeDtypeStruct specs for GPU kernels.
+
+    Returns (aligned_specs, is_scalar_output) where is_scalar_output[i] is True
+    when the i-th output is scalar and should not be padded/unpadded.
+    """
+    import jax  # pyrefly: ignore [import-error, missing-import]
+
+    aligned_specs = []
+    is_scalar = []
+    for shape, dtype in zip(out_shapes, out_dtypes):
+        numel = math.prod(shape)
+        if numel <= 1:
+            aligned_specs.append(jax.ShapeDtypeStruct(shape, dtype))
+            is_scalar.append(True)
+        else:
+            aligned_numel = ((numel + alignment - 1) // alignment) * alignment
+            aligned_specs.append(jax.ShapeDtypeStruct((aligned_numel,), dtype))
+            is_scalar.append(False)
+    return tuple(aligned_specs), is_scalar
+
+
+def pallas_gpu_unpad_results(
+    results: Any,
+    orig_shapes: tuple[Any, ...],
+    is_scalar_output: list[bool] | None = None,
+) -> Any:
+    """Remove padding from GPU kernel results and reshape to original shapes.
+
+    If is_scalar_output is None, all outputs are treated as non-scalar.
+    """
+    if not isinstance(results, tuple):
+        results = (results,)
+    unpadded = []
+    for i, (res, shape) in enumerate(zip(results, orig_shapes)):
+        if is_scalar_output is not None and is_scalar_output[i]:
+            unpadded.append(res)
+        else:
+            orig_numel = math.prod(shape)
+            unpadded.append(res[:orig_numel].reshape(shape))
+    return unpadded[0] if len(unpadded) == 1 else tuple(unpadded)
+
+
+# ---------------------------------------------------------------------------
+# Pallas CPU / TPU tiling helpers
+# ---------------------------------------------------------------------------
+# TPU alignment: last dim must be full or a multiple of 128,
+#                second-to-last dim must be full or a multiple of 8.
+_TPU_ALIGN_LAST = 128
+_TPU_ALIGN_SECOND_LAST = 8
+
+
+def _pallas_tile_size(
+    dim: int, alignment: int, max_tile: int = 1024, is_tpu: bool = False
+) -> int:
+    """Pick the largest aligned tile size <= max_tile for *dim*.
+
+    If *dim* is already <= alignment the full dimension is used (no tiling
+    on this axis).
+    """
+    if dim == 0:
+        # Tile size >= 1 avoids division by zero in JAX's _pad_to_block_dimension
+        return alignment if is_tpu else 1
+
+    if is_tpu:
+        # On TPU, Mosaic requires block dimensions to perfectly align to hardware
+        # registers (128 for inner, 8 for outer). We MUST pad the block spec up to
+        # the next alignment boundary (Mosaic handles the OOB masking).
+        if dim <= alignment:
+            return alignment
+        t = min(max_tile, dim)
+        # Use ceiling division to ensure the tile covers the dimension or aligns upward
+        t = ((t + alignment - 1) // alignment) * alignment
+        return t
+
+    if dim <= alignment:
+        return dim
+    t = min(max_tile, dim)
+    t = (t // alignment) * alignment
+    return max(alignment, t)
+
+
+def pallas_permute(x, perm):
+    """Permute array x according to perm, working around Mosaic TPU bugs.
+
+    Mosaic's jnp.permute_dims produces corrupted output for certain 3D+
+    permutations at large sizes.  Decomposes into: loop over smallest output
+    dim via unrolled Python loop + stack, with recursive sub-permutation.
+    At the 2D base case, uses jnp.permute_dims (which works on Mosaic).
+
+    Callers should ensure tiles are small enough via tiling so that the
+    Mosaic bug doesn't trigger (threshold ~ 256K elements).
+    """
+    import jax.numpy as jnp  # pyrefly: ignore [import-error, missing-import]
+
+    ndim = len(perm)
+    if ndim <= 2:
+        return jnp.permute_dims(x, perm)
+    if perm == tuple(range(ndim)):
+        return x  # identity
+
+    shape = x.shape
+    out_shape = tuple(shape[p] for p in perm)
+
+    # Choose the output dimension with the smallest size to loop over.
+    loop_out_dim = min(range(ndim), key=lambda d: out_shape[d])
+    loop_size = out_shape[loop_out_dim]
+    loop_in_dim = perm[loop_out_dim]
+
+    # Build the sub-permutation for the remaining (ndim-1) dimensions.
+    remaining_in_dims = [d for d in range(ndim) if d != loop_in_dim]
+    remaining_out_perm_raw = [perm[d] for d in range(ndim) if d != loop_out_dim]
+    dim_map = {old: new for new, old in enumerate(remaining_in_dims)}
+    sub_perm = tuple(dim_map[d] for d in remaining_out_perm_raw)
+
+    # Unrolled loop: extract slices with static indices, apply sub-perm,
+    # then stack along the loop output dimension.
+    slices = []
+    # pyrefly: ignore [bad-argument-type]
+    for i in range(loop_size):
+        idx: list[Any] = [slice(None)] * ndim
+        idx[loop_in_dim] = i
+        slc = x[tuple(idx)]  # static index, removes loop_in_dim
+        slc_p = pallas_permute(slc, sub_perm)
+        slices.append(slc_p)
+
+    return jnp.stack(slices, axis=loop_out_dim)
+
+
+# Maximum tile elements for permutation kernels.  Mosaic's
+# jnp.permute_dims produces corrupted output above ~1.5M elements for
+# certain 3D+ permutations.  Keep tiles well under that threshold.
+_PERMUTE_MAX_TILE_ELEMS = 1 << 18  # ~256K elements
+
+
+def pallas_compute_tiling(
+    ref_shape: tuple[int, ...],
+    transpose: bool = False,  # constrain tile size for pallas_permute VMEM
+    skip_last_n: int = 0,
+    exact_only: bool = False,
+    is_tpu: bool = False,
+    permutations: list[tuple[int, ...]] | None = None,
+    max_grid_product: int | None = None,
+) -> tuple[tuple[int, ...], tuple[int, ...], dict[int, int]]:
+    """Compute tile shape, grid and axis→grid-dim mapping for CPU/TPU.
+
+    Always uses TPU-compatible alignment (last dim multiple of 128,
+    second-to-last multiple of 8) so that the same generated kernel works
+    on both CPU-interpret and real TPU.
+
+    *transpose* constrains tile sizes to keep VMEM usage within safe
+    limits for pallas_permute (gather-based permutation).
+
+    *skip_last_n* prevents tiling the last N dimensions (used when those
+    dims correspond to internal reduction ranges that must remain full).
+
+    *exact_only* restricts tiling to dimensions that divide evenly by the
+    tile size (no remainder blocks).  Required on TPU where Mosaic needs
+    block shapes to match the XLA memory layout.
+
+    Returns ``(tile_shape, grid, axis_to_grid)`` where *axis_to_grid*
+    maps each tiled reference-shape axis index to its position in the
+    grid tuple.
+
+    When no dimension benefits from tiling the grid is ``(1,)`` and the
+    tile covers the full tensor.
+    """
+    nd = len(ref_shape)
+    if nd == 0:
+        return (), (1,), {}
+
+    # Effective number of dims eligible for tiling
+    tileable_nd = nd - skip_last_n
+
+    tile = list(ref_shape)
+    grid_parts: list[int] = []
+    axis_to_grid: dict[int, int] = {}  # ref axis → grid dim
+
+    # Pick alignment based on the physical position of the axis in the
+    # tensor, not its position in the tileable subset.  The TPU requires
+    # the physical last dim to be a multiple of 128 and the physical
+    # second-to-last dim to be a multiple of 8.
+    #
+    # When permutations are present, an output dim `ax` maps to input
+    # dim `perm[ax]`.  If `perm[ax]` is the input's last dim, that
+    # output dim must also satisfy _TPU_ALIGN_LAST.  Compute the
+    # strictest alignment across all buffers.
+    def _input_align(input_ax: int, input_nd: int) -> int:
+        return _TPU_ALIGN_LAST if input_ax == input_nd - 1 else _TPU_ALIGN_SECOND_LAST
+
+    def _align(ax: int) -> int:
+        out_align = _TPU_ALIGN_LAST if ax == nd - 1 else _TPU_ALIGN_SECOND_LAST
+        if permutations:
+            for perm in permutations:
+                if perm is not None and len(perm) == nd:
+                    in_ax = perm[ax]
+                    in_align = _input_align(in_ax, len(perm))
+                    out_align = max(out_align, in_align)
+        return out_align
+
+    def _can_tile_ax(ax: int, dim: int, t: int) -> bool:
+        """Check if tiling dim to t is valid."""
+        if t >= dim:
+            # For TPU padding, we allow tiles >= dimension for the aligned axes
+            if _align(ax) == _TPU_ALIGN_LAST or _align(ax) == _TPU_ALIGN_SECOND_LAST:
+                return True
+            return False
+        if exact_only and dim % t != 0:
+            if _align(ax) == _TPU_ALIGN_LAST or _align(ax) == _TPU_ALIGN_SECOND_LAST:
+                # TPU DMA `#tpu.element_window` natively masks out-of-bounds remainder tiles
+                return True
+            return False
+        return True
+
+    # Second-to-last tileable dim (added first so it becomes grid dim 0)
+    if tileable_nd >= 2:
+        ax = tileable_nd - 2
+        t = _pallas_tile_size(ref_shape[ax], _align(ax), is_tpu=is_tpu)
+        if _can_tile_ax(ax, ref_shape[ax], t):
+            tile[ax] = t
+            axis_to_grid[ax] = len(grid_parts)
+            grid_parts.append((ref_shape[ax] + t - 1) // t)
+
+    # Last tileable dim
+    if tileable_nd >= 1:
+        ax = tileable_nd - 1
+        t = _pallas_tile_size(ref_shape[ax], _align(ax), is_tpu=is_tpu)
+        if _can_tile_ax(ax, ref_shape[ax], t):
+            tile[ax] = t
+            axis_to_grid[ax] = len(grid_parts)
+            grid_parts.append((ref_shape[ax] + t - 1) // t)
+
+    # When transpose is active, tile dimensions to keep the total tile
+    # elements below the safe threshold for Mosaic's permute_dims.
+    # Mosaic pads the last dimension to _TPU_ALIGN_LAST (128 for f32),
+    # so when the last dim is small (e.g. 2), VMEM usage is much higher
+    # than the logical tile size.  We compute an effective element count
+    # that accounts for this padding.
+    if transpose:
+        tile_elems = 1
+        for t in tile:
+            tile_elems *= t
+        last_dim = tile[nd - 1] if nd > 0 else 1
+        padded_last = max(last_dim, _TPU_ALIGN_LAST)
+        # pallas_permute decomposes ndim>2 permutations by looping over
+        # the smallest dim and recursing, creating (ndim-1)-D intermediates
+        # at each level.  Mosaic pads the last physical dim of each
+        # intermediate to _TPU_ALIGN_LAST (128), which can cause VMEM OOM
+        # when a small dim (e.g. 4) ends up as the last dim.
+        #
+        # For ndim >= 3, compute effective_max from the worst-case VMEM
+        # of the (ndim-1)-D intermediate created by pallas_permute.
+        # The intermediate removes the loop dim (smallest output dim) and
+        # may have a small dim as its last physical dim -> padded to 128.
+        if nd >= 4:
+            # For 4D+, pallas_permute creates (nd-1)-D intermediates
+            # that can have small dims as their last physical dim,
+            # which Mosaic pads to _TPU_ALIGN_LAST (128).  This can
+            # cause VMEM OOM.  (3D pallas_permute only creates 2D
+            # intermediates which Mosaic handles efficiently.)
+            sorted_dims = sorted(tile)
+            # The loop dim is the smallest; intermediate has the rest.
+            inter_dims = sorted_dims[1:]  # (nd-1) dims, sorted
+            smallest_inter = inter_dims[0]
+            # Padded intermediate: smallest remaining dim padded to 128.
+            padded_inter_elems = 1
+            for d in inter_dims:
+                padded_inter_elems *= d
+            pad_factor = max(smallest_inter, _TPU_ALIGN_LAST) // max(smallest_inter, 1)
+            padded_inter_elems *= pad_factor
+            padded_inter_bytes = padded_inter_elems * 4
+            # Conservative VMEM budget: 2MB per intermediate buffer.
+            # pallas_permute creates multiple live intermediates (input
+            # slices, sub-permutation results, stack temps).  Empirically,
+            # 2MB per worst-case intermediate avoids OOM on TPU v4 (64MB VMEM).
+            vmem_per_buf = 2 * 1024 * 1024  # 2MB
+            if padded_inter_bytes > vmem_per_buf:
+                scale = vmem_per_buf / padded_inter_bytes
+                effective_max = max(int(tile_elems * scale), 1024)
+            else:
+                effective_max = _PERMUTE_MAX_TILE_ELEMS
+        else:
+            # For ndim <= 3: pallas_permute creates 2D intermediates
+            # (fine) but jnp.stack creates a 3D result with the same
+            # shape as the output.  If the last dim is small, Mosaic
+            # pads it to _TPU_ALIGN_LAST (128), which can exhaust VMEM.
+            # Empirically: 8MB per padded buffer keeps us safe on TPU v4
+            # (64MB VMEM, multiple buffers live simultaneously).
+            vmem_per_buf = 8 * 1024 * 1024  # 8MB
+            max_nonlast_product = vmem_per_buf // (padded_last * 4)
+            effective_max = min(
+                max_nonlast_product * last_dim,
+                _PERMUTE_MAX_TILE_ELEMS,
+            )
+        effective_max = min(effective_max, _PERMUTE_MAX_TILE_ELEMS)
+        effective_max = max(effective_max, 1024)
+
+        if tile_elems > effective_max:
+            # Collect tileable dims.  For ndim >= 4 with VMEM-constrained
+            # intermediates, include ALL tileable dims (even the last)
+            # since we may need to shrink every dim to fit in VMEM.
+            # For ndim <= 3, exclude the last dim (it's often small
+            # and already at minimum).
+            if nd >= 4:
+                tileable_axes = [
+                    ax for ax in range(tileable_nd - 1, -1, -1) if ref_shape[ax] > 1
+                ]
+            else:
+                tileable_axes = [
+                    ax for ax in range(tileable_nd - 2, -1, -1) if ref_shape[ax] > 1
+                ]
+            # Distribute reduction proportionally across dims.
+            reduction = tile_elems / effective_max
+            n = len(tileable_axes)
+            per_dim = reduction ** (1.0 / max(n, 1))
+
+            for ax in tileable_axes:
+                if tile_elems <= effective_max:
+                    break
+                dim = ref_shape[ax]
+                cur_t = tile[ax]
+                align = _align(ax)
+                # Target tile: proportional reduction, but also check
+                # against the remaining budget for subsequent dims.
+                target = int(cur_t / per_dim)
+                remaining_budget = effective_max * cur_t // tile_elems
+                target = min(target, remaining_budget)
+                target = (target // align) * align
+                target = max(target, align)
+                # Find largest aligned divisor of dim <= target.
+                best_t = cur_t
+                for candidate in range(target, 0, -align):
+                    if candidate >= cur_t:
+                        continue
+                    if dim % candidate != 0:
+                        continue
+                    best_t = candidate
+                    break
+                if best_t < cur_t:
+                    if ax in axis_to_grid:
+                        grid_parts[axis_to_grid[ax]] = dim // best_t
+                    else:
+                        axis_to_grid[ax] = len(grid_parts)
+                        grid_parts.append(dim // best_t)
+                    tile_elems = tile_elems * best_t // cur_t
+                    tile[ax] = best_t
+
+    grid = tuple(grid_parts) if grid_parts else (1,)
+
+    # Enforce max_grid_product by scaling up tiles on tiled axes.
+    if max_grid_product is not None and grid_parts:
+        grid_product = 1
+        for g in grid_parts:
+            grid_product *= g
+        if grid_product > max_grid_product:
+            # Sort tiled axes by their grid contribution (descending)
+            # and increase tiles to reduce the grid.
+            tiled_axes = sorted(
+                axis_to_grid.keys(),
+                key=lambda ax: grid_parts[axis_to_grid[ax]],
+                reverse=True,
+            )
+            for ax in tiled_axes:
+                if grid_product <= max_grid_product:
+                    break
+                gi = axis_to_grid[ax]
+                cur_grid = grid_parts[gi]
+                cur_tile = tile[ax]
+                dim = ref_shape[ax]
+                align = _align(ax)
+                # Target grid for this axis: proportional share of reduction
+                target_grid = max(1, int(cur_grid * max_grid_product / grid_product))
+                target_tile = (dim + target_grid - 1) // target_grid
+                # Round up to alignment
+                target_tile = ((target_tile + align - 1) // align) * align
+                target_tile = min(target_tile, dim)
+                # Find smallest aligned value >= target_tile that divides dim
+                # (or just use the rounded-up value if exact_only is False)
+                best_tile = dim  # fallback: full dim
+                if exact_only:
+                    for candidate in range(target_tile, dim + 1, align):
+                        if dim % candidate == 0:
+                            best_tile = candidate
+                            break
+                else:
+                    best_tile = target_tile
+                if best_tile > cur_tile:
+                    new_grid = (dim + best_tile - 1) // best_tile
+                    grid_product = grid_product * new_grid // cur_grid
+                    grid_parts[gi] = new_grid
+                    tile[ax] = best_tile
+            grid = tuple(grid_parts)
+
+    return tuple(tile), grid, axis_to_grid
+
+
+def pallas_make_block_spec(
+    buf_shape: tuple[int, ...],
+    ref_shape: tuple[int, ...],
+    tile_shape: tuple[int, ...],
+    axis_to_grid: dict[int, int],
+    n_grid: int,
+    permutation: tuple[int, ...] | None = None,
+    is_output: bool = False,
+) -> Any:
+    """Build a ``pl.BlockSpec`` for *buf_shape* given tiling of *ref_shape*.
+
+    Lower-ndim buffers are right-aligned with the reference shape (numpy
+    broadcast rules).  Dimensions that match a tiled reference dimension
+    are tiled; broadcast dimensions (size 1 or absent) are kept full.
+
+    When *buf_nd > ref_nd* (reduction inputs), we find an alignment offset
+    so the ref dims map into the buffer.  Extra dims are kept at full size
+    with index 0 in the index_map.
+
+    When *permutation* is given (a tuple mapping ref axis -> buf axis),
+    the buffer dimensions are accessed in permuted order relative to the
+    reference shape.  For example, permutation=(1, 0) swaps the last two.
+
+    When *is_output* is True and *buf_nd < ref_nd*, left-alignment is used
+    as a fallback (for reduction outputs whose trailing dims were reduced).
+    """
+    from jax.experimental import (  # pyrefly: ignore [import-error, missing-import]
+        pallas as pl,
+    )
+
+    buf_nd = len(buf_shape)
+    ref_nd = len(ref_shape)
+
+    if buf_nd == 0:
+        # Scalar — untouched regardless of grid shape.
+        return pl.BlockSpec((1,), _make_index_map([], 1, n_grid))
+
+    bs = list(buf_shape)
+    tiled_pairs: list[tuple[int, int]] = []
+
+    if buf_nd > ref_nd:
+        # Reduction input: find alignment offset k where ref dims map into buf.
+        align_k = 0
+        for k in range(buf_nd - ref_nd + 1):
+            ok = True
+            for i in range(ref_nd):
+                if ref_shape[i] == 1:
+                    continue
+                if buf_shape[k + i] != ref_shape[i]:
+                    ok = False
+                    break
+            if ok:
+                align_k = k
+                break
+
+        for ref_ax, grid_dim in axis_to_grid.items():
+            buf_ax = align_k + ref_ax
+            if 0 <= buf_ax < buf_nd and buf_shape[buf_ax] == ref_shape[ref_ax]:
+                bs[buf_ax] = tile_shape[ref_ax]
+                tiled_pairs.append((buf_ax, grid_dim))
+
+    elif permutation is not None and buf_nd == ref_nd:
+        # Permuted buffer: map ref axes through the permutation.
+        for ref_ax, grid_dim in axis_to_grid.items():
+            buf_ax = permutation[ref_ax]
+            if 0 <= buf_ax < buf_nd and buf_shape[buf_ax] == ref_shape[ref_ax]:
+                bs[buf_ax] = tile_shape[ref_ax]
+                tiled_pairs.append((buf_ax, grid_dim))
+
+    else:
+        # Standard right-alignment, with left-alignment fallback for
+        # reduction outputs (e.g. sum(dim=-1) on (10,10) → (10,)).
+        for ref_ax, grid_dim in axis_to_grid.items():
+            buf_ax = ref_ax - (ref_nd - buf_nd)
+            if 0 <= buf_ax < buf_nd and buf_shape[buf_ax] == ref_shape[ref_ax]:
+                bs[buf_ax] = tile_shape[ref_ax]
+                tiled_pairs.append((buf_ax, grid_dim))
+            elif (
+                is_output
+                and buf_nd < ref_nd
+                and 0 <= ref_ax < buf_nd
+                and buf_shape[ref_ax] == ref_shape[ref_ax]
+            ):
+                # Left-aligned match for output: buf dim i matches ref dim i
+                # (reduction output whose trailing dims were reduced away)
+                bs[ref_ax] = tile_shape[ref_ax]
+                tiled_pairs.append((ref_ax, grid_dim))
+
+    return pl.BlockSpec(
+        tuple(bs),
+        _make_index_map(tiled_pairs, buf_nd, n_grid),
+    )
+
+
+def _make_index_map(
+    tiled_pairs: list[tuple[int, int]],
+    buf_nd: int,
+    n_grid: int,
+) -> Any:
+    """Return an index_map callable for ``pl.BlockSpec``.
+
+    *tiled_pairs* is a list of ``(buf_axis, grid_dim)`` indicating which
+    buffer axes receive a grid index.  All other axes return 0 (full block).
+
+    All returned values are explicitly ``jnp.int32`` so that TPU Mosaic
+    lowering (which rejects 64-bit types) works when ``jax_enable_x64`` is
+    active.  The casts are created inside the function body (not captured)
+    to satisfy JAX's "index_map must not capture constants" rule.
+    """
+    import jax.numpy as jnp  # pyrefly: ignore [import-error, missing-import]
+
+    # Pre-build the mapping so the returned lambda is a plain lookup.
+    mapping = dict(tiled_pairs)
+
+    if n_grid == 0 or (n_grid == 1 and not mapping):
+        return lambda _i: tuple(jnp.int32(0) for _ in range(buf_nd))
+
+    def index_map(*grid_args):
+        return tuple(
+            jnp.int32(grid_args[mapping[d]]) if d in mapping else jnp.int32(0)
+            for d in range(buf_nd)
+        )
+
+    return index_map
+
+
+def pallas_ensure_nonzero_rank(x: torch.Tensor) -> torch.Tensor:
+    if len(x.shape) == 0:
+        return x.reshape((1,))
+    return x
+
+
+def pallas_make_block_spec_non_tiled(shape: tuple[int, ...]) -> Any:
+    import jax.numpy as jnp  # pyrefly: ignore [import-error, missing-import]
+    from jax.experimental import (  # pyrefly: ignore [import-error, missing-import]
+        pallas as pl,
+    )
+
+    nonzero_rank_shape = shape if len(shape) > 0 else (1,)
+    return pl.BlockSpec(
+        nonzero_rank_shape,
+        lambda i: [jnp.int32(i)] * len(nonzero_rank_shape),
+    )

outputs/audit_venv/lib/python3.11/site-packages/torch/_inductor/runtime/static_triton_launcher.py

@@ -0,0 +1,351 @@
+import functools
+import os
+from functools import cached_property
+from typing import Any
+from typing_extensions import Unpack
+
+from ..utils import is_rocm
+from .triton_compat import ASTSource, CompiledKernel, knobs as triton_knobs
+from .triton_helpers import get_constexprs
+
+
+class StaticallyLaunchedTritonKernel:
+    """
+    Parses the metadata of a CompiledKernel from Triton into a structure that can
+    launch the cuda kernel directly. Only works for triton kernels compiled to cubin.
+
+    Doing this avoids C++ codegen and compilation during compile, since we can use a
+    statically compiled library to launch the kernel. To avoid mallocing for the arguments,
+    we have a launcher for different numbers of arguments up to a max. StaticCudaLauncher
+    only supports # of arguments up until 10 for now.
+
+    Workflow:
+    Compile time:
+    1. Compile a kernel with triton and get a CompiledKernel
+    2. Instantiate kernel = StaticallyLaunchedTritonKernel(triton_kernel)
+    3. Write to a cubin file: kernel.write_cubin_to_file(filepath)
+    4. Call kernel.load_kernel() (CUDA should be initialized by this point) to load the cubin
+    Runtime:
+    5. Call kernel.run(grid, stream, args) to launch the kernel
+
+    Note that after step 3, StaticallyLaunchedTritonKernel is fully pickleable/serializable.
+    This allows it to be cached by FXGraphCache/TritonBundler, as well as sent from the worker
+    to the parent process in inductor.
+
+    There are two main versions of triton that we wish to support: 3.3 and 3.2. Triton makes considerable changes
+    to how it handles constants in 3.3, so there's some special logic necessary to handle both versions.
+    """
+
+    @cached_property
+    def C_impl(self):
+        raise NotImplementedError
+
+    def __init__(self, kernel: CompiledKernel) -> None:
+        # pyrefly: ignore [missing-attribute]
+        self.name = kernel.src.fn.__name__
+        # pyrefly: ignore [missing-attribute]
+        self.cubin_path = kernel._cubin_path
+
+        # Used by torch.compile to filter constants in older triton versions
+        # pyrefly: ignore [missing-attribute]
+        self.arg_names = kernel.src.fn.arg_names
+
+        # Const exprs that are declared by the triton kernel directly
+        # Used to generate the kernel launcher's def args
+        # pyrefly: ignore [missing-attribute]
+        self.declared_constexprs = get_constexprs(kernel.src.fn)
+
+        # pyrefly: ignore [missing-attribute]
+        self.hash = kernel.hash
+
+        if triton_knobs is None:
+            # pyrefly: ignore [missing-attribute]
+            launch_enter = kernel.__class__.launch_enter_hook
+            # pyrefly: ignore [missing-attribute]
+            launch_exit = kernel.__class__.launch_exit_hook
+        else:
+            launch_enter = triton_knobs.runtime.launch_enter_hook
+            launch_exit = triton_knobs.runtime.launch_exit_hook
+
+        def hook_is_empty(hook: Any) -> bool:
+            if hook is None:
+                return True
+            if (
+                triton_knobs
+                and (HookChain := getattr(triton_knobs, "HookChain", None)) is not None
+                and isinstance(hook, HookChain)
+            ):
+                # Support hooks after https://github.com/triton-lang/triton/pull/7866
+                return len(hook.calls) == 0
+            return False
+
+        if not hook_is_empty(launch_enter) or not hook_is_empty(launch_exit):
+            raise NotImplementedError(
+                "We don't support launch enter or launch exit hooks"
+            )
+        # pyrefly: ignore [missing-attribute]
+        self.num_warps = kernel.metadata.num_warps
+        self.shared = (
+            # pyrefly: ignore [missing-attribute]
+            kernel.shared if hasattr(kernel, "shared") else kernel.metadata.shared
+        )
+
+        def needs_scratch_arg(scratch_name: str, param_name: str) -> bool:
+            # pyrefly: ignore [missing-attribute]
+            if hasattr(kernel.metadata, param_name):
+                # pyrefly: ignore [missing-attribute]
+                if getattr(kernel.metadata, param_name) > 0:
+                    raise NotImplementedError(
+                        f"{scratch_name} scratch not yet supported"
+                    )
+                return True
+            return False
+
+        # Newer triton versions pass an extra global scratch parameter to the compiled cuda kernel.
+        # Inductor never uses this field or enables it, but we still have to pass
+        # an extra None into the set of params if its enabled
+        self.has_global_scratch = needs_scratch_arg("Global", "global_scratch_size")
+        # same situation for profile scratch - triton-lang/triton#7258
+        self.has_profile_scratch = needs_scratch_arg("Profile", "profile_scratch_size")
+
+        # pyrefly: ignore [missing-attribute]
+        self.arg_tys = self.arg_ty_from_signature(kernel.src)
+        self.function: int | None = None  # Loaded by load_kernel(on the parent process)
+        num_ctas = 1
+        if hasattr(kernel, "num_ctas"):
+            num_ctas = kernel.num_ctas
+        elif hasattr(kernel, "metadata"):
+            num_ctas = kernel.metadata.num_ctas
+
+        if num_ctas != 1:
+            raise NotImplementedError(
+                "Static cuda launcher only supports num_ctas == 1"
+            )
+
+    def reload_cubin_from_raw(self, filepath: str) -> str:
+        """
+        If the cubin file triton generated gets deleted under us, we can
+        reload it from the raw cubin file.
+        """
+        if self.cubin_path is None:
+            assert self.cubin_raw is not None
+            os.makedirs(os.path.dirname(filepath), exist_ok=True)
+            with open(filepath, "wb") as f:
+                f.write(self.cubin_raw)
+                self.cubin_path = filepath  # pyre-ignore
+        return self.cubin_path
+
+    def load_kernel(self, device: int) -> None:
+        if self.function is not None:
+            return
+
+        assert hasattr(self, "cubin_path")
+        assert self.cubin_path is not None
+        (self.function, self.n_regs, self.n_spills) = self.C_impl._load_kernel(
+            self.cubin_path, self.name, self.shared, device
+        )
+        # Don't need the cubin path anymore now that we've loaded
+        self.cubin_path = None
+        self.cubin_raw = None
+
+    @staticmethod
+    @functools.lru_cache
+    def type_mappings() -> dict[str, str]:
+        return {
+            "i1": "i",
+            "i8": "b",
+            "i16": "h",
+            "i32": "i",
+            "i64": "l",
+            "u1": "I",
+            "u8": "B",
+            "u16": "H",
+            "u32": "I",
+            "u64": "K",
+            "fp16": "f",
+            "bf16": "f",
+            "fp32": "f",
+            "f32": "f",
+            "fp64": "d",
+            # TODO handle nvTmaDesc/CUtensormap
+        }
+
+    def extract_type(self, ty: str) -> str:
+        """
+        Takes a triton type from CompiledKernel.signature and
+        converts it into a single char encoding. _StaticCudaLauncher
+        will switch on this char to figure out what type the underlying
+        value should be passed to the triton kernel as.
+        """
+        if ty[0] == "*":
+            return "O"
+        elif ty == "nvTmaDesc":
+            raise NotImplementedError("nvTmaDesc kernels are not yet supported")
+        return StaticallyLaunchedTritonKernel.type_mappings()[ty]
+
+    def arg_ty_from_signature(self, src: ASTSource) -> str:
+        def index_key(i: Any) -> int:
+            if isinstance(i, str):
+                # pyrefly: ignore [missing-attribute]
+                return src.fn.arg_names.index(i)
+            elif isinstance(i, tuple):
+                # In triton 3.3, src.fn.constants has tuples as a key
+                return i[0]
+            else:
+                return i
+
+        # pyrefly: ignore [missing-attribute]
+        signature = {index_key(key): value for key, value in src.signature.items()}
+        # Triton uses these as the main way to filter out constants passed to their cubin
+        constants = [index_key(key) for key in getattr(src, "constants", dict())]
+        # This value is always a superset of kernel.fn.constexprs: kernel.fn.constexprs are
+        # constants declared by the triton kernel directly, whereas this list can have
+        # constants that are unused by the triton kernel that triton figured out during
+        # compilation.
+        self.full_constexprs = constants
+        # Despite requiring them to be passed in, the triton CUDA launcher
+        # completely ignores the constexprs passed into it when generating code.
+        # So we can ignore them here too
+        params = []
+
+        for i in sorted(signature.keys()):
+            ty = signature[i]
+            # In newer triton versions, constants are passed in to signature with type `constexpr`
+            # In older triton versions, there can be constants in src.constants that are not `constexpr` in signature
+            # so we check both here
+            if ty == "constexpr" or i in constants:
+                pass
+            else:
+                # pyrefly: ignore [bad-argument-type]
+                params.append(self.extract_type(ty))
+        return "".join(params)
+
+    def __getstate__(self) -> dict[str, Any]:
+        # Remove objects that are no longer valid for pickling
+        state = self.__dict__.copy()
+        state["function"] = None
+        # Cubin paths aren't consistent across processes, so we clear
+        # and reload them.
+        state["cubin_path"] = None
+        return state
+
+    def run(
+        self,
+        grid_x: int,
+        grid_y: int,
+        grid_z: int,
+        stream: int,
+        *args: Unpack[tuple[object, ...]],
+    ) -> None:
+        """Actually run the kernel at runtime. This function is the hot codepath."""
+
+        # Assert load_kernel() has been called and args match
+        assert self.function is not None
+
+        # TODO: actually, if the args *don't* match, we probably should
+        # throw an exception. But if inductor is the only one calling this
+        # thing, it should always match.
+        # Get rid of constants before passing to cubin launcher
+
+        arg_tys = self.arg_tys
+
+        if is_rocm():
+            # ROCm/HIP kernel ABI: The Triton HIP backend ALWAYS includes both
+            # global_scratch and profile_scratch parameters in the kernel signature,
+            # even when the kernel doesn't use them (i.e., when has_*_scratch is False).
+            #
+            # This differs fundamentally from CUDA, where these parameters are only
+            # present in the signature if the corresponding has_*_scratch flag is True.
+            #
+            # The flags indicate whether memory will be allocated/used:
+            # - has_global_scratch: Whether global scratch workspace is needed
+            # - has_profile_scratch: Whether profiling instrumentation is enabled
+            #
+            # However, regardless of flag values, we MUST always pass both parameters
+            # to match the HIP kernel ABI. Passing None is safe:
+            #
+            # - If scratch is not needed (has_*_scratch=False or scratch_size=0):
+            #   The None becomes nullptr, which the kernel never dereferences
+            #
+            # - If scratch is needed (has_*_scratch=True and scratch_size>0):
+            #   The None becomes nullptr initially, but the HIP runtime intercepts
+            #   the kernel launch, allocates the required scratch memory based on
+            #   kernel metadata, and replaces the nullptr with a valid pointer before
+            #   the kernel actually executes
+            #
+            # Not passing both parameters causes segmentation faults because the kernel
+            # expects them at specific positions in the argument array.
+            arg_tys = arg_tys + "OO"
+            args = (*args, None, None)
+
+        else:
+            for has_scratch in [self.has_global_scratch, self.has_profile_scratch]:
+                if has_scratch:
+                    arg_tys = arg_tys + "O"
+                    args = (*args, None)
+        # pyrefly: ignore [bad-argument-type]
+        assert len(args) == len(arg_tys)
+
+        # TODO: can handle grid functions here or in C++, so
+        # that we don't need the grid handler above.
+        self.C_impl._launch_kernel(
+            self.function,
+            grid_x,
+            grid_y,
+            grid_z,
+            self.num_warps,
+            self.shared,
+            arg_tys,
+            args,
+            stream,
+        )
+
+
+class StaticallyLaunchedCudaKernel(StaticallyLaunchedTritonKernel):
+    @cached_property
+    def C_impl(self):
+        from torch._C import _StaticCudaLauncher
+
+        return _StaticCudaLauncher
+
+    def __init__(self, kernel: CompiledKernel) -> None:
+        # pyrefly: ignore [missing-attribute]
+        if "hsaco" in kernel.asm:
+            # pyrefly: ignore [missing-attribute]
+            self.cubin_raw = kernel.asm["hsaco"]
+
+        # pyrefly: ignore [missing-attribute]
+        elif "cubin" in kernel.asm:
+            # pyrefly: ignore [missing-attribute]
+            self.cubin_raw = kernel.asm["cubin"]
+        else:
+            raise RuntimeError(
+                "Expected either 'hsaco' (ROCm) or 'cubin' (CUDA) in kernel.asm"
+            )
+        super().__init__(kernel)
+
+
+class StaticallyLaunchedXpuKernel(StaticallyLaunchedTritonKernel):
+    @cached_property
+    def C_impl(self):
+        from torch._C import _StaticXpuLauncher
+
+        return _StaticXpuLauncher
+
+    def __init__(self, kernel: CompiledKernel) -> None:
+        # pyrefly: ignore [missing-attribute]
+        self.cubin_raw = kernel.asm.get("zebin", None)
+        super().__init__(kernel)
+
+
+def statically_launched_kernel_by_device(
+    kernel: CompiledKernel, device_type: str = "cuda"
+) -> StaticallyLaunchedTritonKernel:
+    if device_type in ("cuda", "hip"):
+        return StaticallyLaunchedCudaKernel(kernel)
+    elif device_type == "xpu":
+        return StaticallyLaunchedXpuKernel(kernel)
+    else:
+        raise NotImplementedError(
+            f"Device type {device_type} is not supported for static launcher"
+        )

outputs/audit_venv/lib/python3.11/site-packages/torch/_inductor/runtime/triton_compat.py

@@ -0,0 +1,186 @@
+from __future__ import annotations
+
+import inspect
+from typing import Any
+
+import torch
+
+
+try:
+    import triton
+except ImportError:
+    triton = None
+
+
+if triton is not None:
+    import triton.language as tl
+    from triton import Config
+    from triton.compiler import CompiledKernel
+    from triton.runtime.autotuner import OutOfResources
+    from triton.runtime.jit import JITFunction, KernelInterface
+
+    try:
+        from triton.runtime.autotuner import PTXASError
+    except ImportError:
+
+        class PTXASError(Exception):  # type: ignore[no-redef]
+            pass
+
+    try:
+        from triton.compiler.compiler import ASTSource
+    except ImportError:
+        ASTSource = None
+
+    try:
+        from triton.backends.compiler import GPUTarget
+    except ImportError:
+
+        def GPUTarget(
+            backend: str,
+            arch: int | str,
+            warp_size: int,
+        ) -> Any:
+            if torch.version.hip:
+                return [backend, arch, warp_size]
+            return (backend, arch)
+
+    # In the latest triton, math functions were shuffled around into different modules:
+    # https://github.com/triton-lang/triton/pull/3172
+    try:
+        from triton.language.extra import libdevice
+
+        libdevice = tl.extra.libdevice  # noqa: F811
+        math = tl.math
+    except ImportError:
+        if hasattr(tl.extra, "cuda") and hasattr(tl.extra.cuda, "libdevice"):
+            libdevice = tl.extra.cuda.libdevice
+            math = tl.math
+        elif hasattr(tl.extra, "intel") and hasattr(tl.extra.intel, "libdevice"):
+            libdevice = tl.extra.intel.libdevice
+            math = tl.math
+        else:
+            libdevice = tl.math
+            math = tl
+
+    try:
+        from triton.language.standard import _log2
+    except ImportError:
+
+        def _log2(x: Any) -> Any:
+            raise NotImplementedError
+
+    def _triton_config_has(param_name: str) -> bool:
+        if not hasattr(triton, "Config"):
+            return False
+        if not hasattr(triton.Config, "__init__"):
+            return False
+        return param_name in inspect.signature(triton.Config.__init__).parameters
+
+    # Drop the legacy support of autoWS
+    HAS_WARP_SPEC = False
+
+    try:
+        from triton import knobs
+    except ImportError:
+        knobs = None
+
+    try:
+        from triton.runtime.cache import triton_key  # type: ignore[attr-defined]
+    except ImportError:
+        from triton.compiler.compiler import (
+            triton_key,  # type: ignore[attr-defined,no-redef]
+        )
+
+    try:
+        from triton.runtime.errors import IntelGPUError
+    except ImportError:
+
+        class IntelGPUError(Exception):  # type: ignore[no-redef]
+            pass
+
+    builtins_use_semantic_kwarg = (
+        "_semantic" in inspect.signature(triton.language.core.view).parameters
+    )
+    HAS_TRITON = True
+else:
+
+    def _raise_error(*args: Any, **kwargs: Any) -> Any:
+        raise RuntimeError("triton package is not installed")
+
+    class OutOfResources(Exception):  # type: ignore[no-redef]
+        pass
+
+    class PTXASError(Exception):  # type: ignore[no-redef]
+        pass
+
+    class IntelGPUError(Exception):  # type: ignore[no-redef]
+        pass
+
+    Config = object
+    CompiledKernel = object
+    KernelInterface = object
+    ASTSource = None
+    GPUTarget = None
+    _log2 = _raise_error
+    libdevice = None
+    math = None
+    knobs = None
+    builtins_use_semantic_kwarg = False
+
+    class triton:  # type: ignore[no-redef]
+        @staticmethod
+        def jit(*args: Any, **kwargs: Any) -> Any:
+            return _raise_error
+
+    class tl:  # type: ignore[no-redef]
+        @staticmethod
+        def constexpr(val: Any) -> Any:
+            return val
+
+        tensor = Any
+        dtype = Any
+
+    class JITFunction:  # type: ignore[no-redef]
+        pass
+
+    HAS_WARP_SPEC = False
+    triton_key = _raise_error
+    HAS_TRITON = False
+
+
+def cc_warp_size(cc: str | int) -> int:
+    if torch.version.hip:
+        if "gfx9" in str(cc):
+            return 64
+        else:
+            return 32
+    else:
+        return 32
+
+
+try:
+    autograd_profiler = torch.autograd.profiler
+except AttributeError:  # Compile workers only have a mock version of torch
+
+    class autograd_profiler:  # type: ignore[no-redef]
+        _is_profiler_enabled = False
+
+
+__all__ = [
+    "Config",
+    "CompiledKernel",
+    "OutOfResources",
+    "KernelInterface",
+    "PTXASError",
+    "IntelGPUError",
+    "ASTSource",
+    "GPUTarget",
+    "tl",
+    "_log2",
+    "libdevice",
+    "math",
+    "triton",
+    "cc_warp_size",
+    "knobs",
+    "triton_key",
+]

outputs/audit_venv/lib/python3.11/site-packages/torch/_inductor/runtime/triton_helpers.py

@@ -0,0 +1,850 @@
+# mypy: allow-untyped-decorators
+# mypy: allow-untyped-defs
+import math as pymath
+import warnings
+from collections.abc import Callable
+from typing import Any, TypeVar
+
+from .triton_compat import (  # noqa: F401
+    _log2,
+    builtins_use_semantic_kwarg,
+    JITFunction,
+    libdevice,
+    math,
+    tl,
+    triton,
+)
+
+
+_T = TypeVar("_T")
+_LOG_2_E: tl.constexpr = tl.constexpr(pymath.log2(pymath.e))
+
+
+def set_driver_to_cpu():
+    driver = triton.runtime.driver
+    if backend := triton.backends.backends.get("cpu", None):
+        if isinstance(driver.active, backend.driver):
+            # Don't re-initialize backend if it is already active
+            return
+        driver.set_active(backend.driver())
+        return
+    # This can be a hard error once triton-cpu is merged into fbcode
+    warnings.warn(
+        "Could not find an active CPU backend. Generated kernels will not be executable!"
+    )
+
+
+def _is_backend_active(name, backend):
+    if backend.driver.is_active():
+        return True
+    # Triton may fail to detect the GPU in subprocess workers when using
+    # ctypes-based driver detection (triton-lang/triton#9578). Fall back
+    # to torch's own device checks which are more reliable in these environments.
+    if name == "nvidia":
+        import torch
+
+        return torch.cuda.is_available() and torch.version.hip is None
+    if name == "amd":
+        import torch
+
+        return torch.cuda.is_available() and torch.version.hip is not None
+    return False
+
+
+def set_driver_to_gpu():
+    driver = triton.runtime.driver
+    for name, backend in triton.backends.backends.items():
+        if _is_backend_active(name, backend) and name != "cpu":
+            # After https://github.com/triton-lang/triton/commit/b844d519bc5e86edf00fe6b3c6c2d1badcd509a4,
+            # `driver.active` can be of `LazyProxy` type and the sign of this - `_obj` attribute.
+            if (
+                isinstance(driver.active, backend.driver)
+                or hasattr(driver.active, "_obj")
+                and isinstance(driver.active._obj, backend.driver)
+            ):
+                # Don't re-initialize backend if it is already active
+                return
+            driver.set_active(backend.driver())
+            return
+    raise RuntimeError("Could not find an active GPU backend")
+
+
+def get_backend_options():
+    from triton.runtime import driver
+
+    target = driver.active.get_current_target()
+    backend = triton.compiler.compiler.make_backend(target)
+    options = backend.parse_options(dict())
+    return options.__dict__
+
+
+def get_constexprs(kernel: JITFunction) -> list[int]:
+    return [p.num for p in kernel.params if p.is_constexpr]
+
+
+@triton.jit
+def promote_to_tensor(x):
+    # Addition promotes to tensor for us
+    return x + tl.zeros((1,), tl.int1)
+
+
+@triton.jit
+def div_floor_integer(a, b):
+    # NOTE: a // b is C division, but we want floor division
+    # Based on c10::div_floor_integer
+    quot = a // b
+    remainder = a % b
+    fixed = tl.where(remainder != 0, quot - 1, quot)
+    return tl.where((a < 0) != (b < 0), fixed, quot)
+
+
+@triton.jit
+def remainder_integer(a, b):
+    # NOTE: a % b matches C division, not floor division
+    remainder = a % b
+    return tl.where((remainder != 0) & ((a < 0) != (b < 0)), remainder + b, remainder)
+
+
+@triton.jit
+def pow_integer(base, exponent):
+    # Triton has no exact integer pow primitive; use repeated squaring for
+    # nonnegative integer exponents so integral scalar pow does not round
+    # through libdevice.pow before casting back to int.
+    exponent_dtype: tl.constexpr = tl.core.get_int_dtype(
+        exponent.dtype.primitive_bitwidth, signed=False
+    )
+    exp = exponent.to(exponent_dtype)
+    result = tl.full(base.shape, 1, base.dtype)
+    for _ in tl.static_range(exponent_dtype.primitive_bitwidth):
+        result = tl.where((exp & 1) != 0, result * base, result)
+        exp = exp >> 1
+        base = base * base
+    return result
+
+
+@triton.jit
+def is_floating(x):
+    return promote_to_tensor(x).dtype.is_floating()
+
+
+@triton.jit
+def _prod_accumulate(a, b):
+    return a * b
+
+
+@triton.jit
+def prod(input, axis):
+    return tl.reduce(input, axis, _prod_accumulate)
+
+
+@triton.jit
+def minimum(a, b):
+    mask = a < b
+    if is_floating(a):
+        mask |= a != a
+    return tl.where(mask, a, b)
+
+
+@triton.jit
+def maximum(a, b):
+    mask = a > b
+    if is_floating(a):
+        mask |= a != a
+    return tl.where(mask, a, b)
+
+
+@triton.jit
+def min2(a, dim):
+    return tl.reduce(a, dim, minimum)
+
+
+@triton.jit
+def max2(a, dim):
+    return tl.reduce(a, dim, maximum)
+
+
+@triton.jit
+def minimum_with_index(a_value, a_index, b_value, b_index):
+    mask = a_value < b_value
+    equal = a_value == b_value
+    if is_floating(a_value):
+        a_isnan = a_value != a_value
+        b_isnan = b_value != b_value
+        mask |= a_isnan & (not b_isnan)
+        # Consider NaNs as equal
+        equal |= a_isnan & b_isnan
+
+    # Prefer lowest index if values are equal
+    mask |= equal & (a_index < b_index)
+    return tl.where(mask, a_value, b_value), tl.where(mask, a_index, b_index)
+
+
+@triton.jit
+def maximum_with_index(a_value, a_index, b_value, b_index):
+    mask = a_value > b_value
+    equal = a_value == b_value
+    if is_floating(a_value):
+        a_isnan = a_value != a_value
+        b_isnan = b_value != b_value
+        mask |= a_isnan & (not b_isnan)
+        # Consider NaNs as equal
+        equal |= a_isnan & b_isnan
+
+    # Prefer lowest index if values are equal
+    mask |= equal & (a_index < b_index)
+    return tl.where(mask, a_value, b_value), tl.where(mask, a_index, b_index)
+
+
+@triton.jit
+def min_with_index(value, index, dim):
+    return tl.reduce((value, index), dim, minimum_with_index)
+
+
+@triton.jit
+def max_with_index(value, index, dim):
+    return tl.reduce((value, index), dim, maximum_with_index)
+
+
+@triton.jit
+def exp(x, use_fast_math: tl.constexpr):
+    if use_fast_math:
+        return math.exp(x)
+    else:
+        return libdevice.exp(x)
+
+
+@triton.jit
+def online_softmax_reduce(lhs_max, lhs_sum, dim, use_fast_math: tl.constexpr):
+    out_max = max2(lhs_max, dim)
+    out_max_keepdim = tl.expand_dims(out_max, dim)
+    delta = tl.where(out_max_keepdim == float("-inf"), 0, lhs_max - out_max_keepdim)
+    out_sum = tl.sum(lhs_sum * exp(delta, use_fast_math), dim)
+    return out_max, out_sum
+
+
+@triton.jit
+def online_softmax_combine(lhs_max, lhs_sum, rhs_max, use_fast_math: tl.constexpr):
+    """
+    When we do combine, we assume lhs is the accumulator and rhs is the next
+    block of data.
+    Then rhs_sum is always 1. With that assumption, we can save some registers
+    and computation.
+    """
+    out_max = maximum(lhs_max, rhs_max)
+
+    lhs_scale = tl.where(
+        out_max == float("-inf"), 1.0, exp(lhs_max - out_max, use_fast_math)
+    )
+    rhs_scale = tl.where(
+        out_max == float("-inf"), 1.0, exp(rhs_max - out_max, use_fast_math)
+    )
+
+    # Should be
+    #   out_sum = lhs_sum * lhs_scale + rhs_sum * rhs_scale
+    # but since rhs_sum is all 1, we can simplify it.
+    out_sum = lhs_sum * lhs_scale + rhs_scale
+    return out_max, out_sum
+
+
+@triton.jit
+def welford_reduce(value, mean, m2, weight, first_iteration):
+    if first_iteration:
+        new_weight = tl.full(weight.shape, 1, weight.dtype)
+        new_mean = value
+        new_m2 = tl.zeros_like(m2)
+    else:
+        delta = value - mean
+        new_weight = weight + 1
+        new_mean = mean + delta / new_weight
+        new_m2 = m2 + delta * (value - new_mean)
+    return new_mean, new_m2, new_weight
+
+
+@triton.jit
+def welford_combine(mean_1, m2_1, weight_1, mean_2, m2_2, weight_2):
+    delta = mean_2 - mean_1
+    new_weight = weight_1 + weight_2
+    w2_over_w = tl.where(new_weight == 0.0, 0.0, weight_2 / new_weight)
+    return (
+        mean_1 + delta * w2_over_w,
+        m2_1 + m2_2 + delta * delta * weight_1 * w2_over_w,
+        new_weight,
+    )
+
+
+@triton.jit
+def welford(mean, m2, weight, dim):
+    return tl.reduce((mean, m2, weight), dim, welford_combine)
+
+
+@triton.jit
+def device_assert_then(cond, msg, r):
+    tl.device_assert(cond, msg)
+    return r
+
+
+@triton.jit
+def rand_eager_kernel(seed, offset_blocks, tid: tl.tensor, VEC: tl.constexpr):
+    inv = 1.0 / 4294967296.0
+    half = inv * 0.5
+
+    tid_u64 = tid.to(tl.uint64)
+
+    subseq = tid_u64 // VEC
+    which4 = (tid_u64 % VEC) // 4
+    lane = tid_u64 % 4
+
+    offblk = offset_blocks.to(tl.uint64) + which4
+
+    u0, u1, u2, u3 = tl.philox(
+        seed,
+        (offblk & 0xFFFFFFFF).to(tl.uint32),
+        ((offblk >> 32) & 0xFFFFFFFF).to(tl.uint32),
+        (subseq & 0xFFFFFFFF).to(tl.uint32),
+        ((subseq >> 32) & 0xFFFFFFFF).to(tl.uint32),
+    )
+
+    v01 = tl.where(lane == 0, u0, u1)
+    v23 = tl.where(lane == 2, u2, u3)
+    rand_int = tl.where((lane == 0) | (lane == 1), v01, v23)
+
+    return 1.0 - (rand_int.to(tl.float32) * inv + half)
+
+
+@triton.jit
+def randint64(seed, offset, low, high):
+    r0, r1, _r2, _r3 = tl.randint4x(seed, offset)
+    r0 = r0.to(tl.uint64)
+    r1 = r1.to(tl.uint64)
+    result = r0 | (r1 << 32)
+    size = high - low
+    result = result % size.to(tl.uint64)
+    result = result.to(tl.int64) + low
+    return result
+
+
+@triton.jit
+def _any_combine(a, b):
+    return a | b
+
+
+@triton.jit
+def any(a, dim):
+    return tl.reduce(a, dim, _any_combine)
+
+
+@triton.jit
+def bucketize_binary_search(
+    values: tl.tensor,
+    boundaries_ptr: tl.tensor,
+    BOUNDARIES_SIZE: int,
+    BOUNDARIES_UNDERLYING_NUMEL: int,
+    BOUNDARIES_STRIDE: int,
+    boundary_indices: tl.tensor,
+    indexing_dtype: tl.dtype,
+    right: "bool",  # triton can't handle the unquoted bool annotation
+    sorter_ptr: tl.tensor,
+    SORTER_STRIDE: int,
+    sorter_indices: tl.tensor,
+):
+    """
+    See [Note: Inductor bucketize op]
+
+    Inputs:
+    -------
+    values: the values to bucketize.
+    boundaries_ptr: a pointer to the beginning of the boundaries tensor, in 1-D.
+    BOUNDARIES_SIZE: the length of the last dimension of the boundaries tensor (i.e. one
+    individual set of boundaries).
+    BOUNDARIES_UNDERLYING_NUMEL: the length of the boundaries tensor, in 1-D, ignoring
+    any striding.
+    BOUNDARIES_STRIDE: the stride of the last dimension of the boundaries tensor
+    boundary_indices: a tensor of the same size as "values"; each element is an index
+    into a 1-D, un-strided boundaries tensor, pointing to the first element in the set
+    of boundaries used for that value.
+    indexing_dtype: the dtype used for indexing into the boundaries tensor, and the
+    return dtype.
+    right: if true, use boundary intervals closed on the left; otherwise use intervals
+    closed on the right.
+    sorter_ptr: an optional pointer to a sorter tensor of the same shape as boundaries,
+    but potentially different striding.  If present, this allows us to treat boundaries
+    as sorted even if the elements of boundaries are unsorted.
+    SORTER_STRIDE: must be present if sorter_ptr is non-None; the stride of the last
+    dimension of the sorter tensor.
+    sorter_indices: must be present if sorter_ptr is non-None; see "boundary_indices".
+    BLOCK_SHAPE: the shape of the data block being processed.
+    """
+
+    low = tl.zeros(values.shape, dtype=indexing_dtype)
+    high = tl.full(values.shape, BOUNDARIES_SIZE, dtype=indexing_dtype)
+
+    full_range = BOUNDARIES_SIZE + 1
+    while full_range > 1:
+        mid = (high + low) // 2
+        mask = (
+            (mid * BOUNDARIES_STRIDE + boundary_indices) < BOUNDARIES_UNDERLYING_NUMEL
+        ).logical_and(mid < BOUNDARIES_SIZE)
+        mid_indices = (
+            mid
+            if sorter_ptr is None or SORTER_STRIDE is None
+            else tl.load(
+                sorter_ptr + sorter_indices + SORTER_STRIDE * mid,
+                mask=mask,
+                other=0,
+            )
+        )
+
+        bucket_upper_bound = tl.load(
+            boundaries_ptr + boundary_indices + BOUNDARIES_STRIDE * mid_indices,
+            mask=mask,
+            other=0,
+        )
+        if right:
+            is_above = values >= bucket_upper_bound
+        else:
+            is_above = values > bucket_upper_bound
+
+        if is_floating(values):
+            is_above = is_above | (values != values)
+
+        low = tl.where(is_above & mask, mid + 1, low)
+        high = tl.where(is_above, high, mid)
+
+        full_range = (full_range + 1) // 2
+
+    return low
+
+
+@triton.jit
+def pack_value_flag(
+    value,
+    flag,
+    DTYPE_VALUE_AS_UINT: tl.constexpr,
+    DTYPE_PACK: tl.constexpr,
+):
+    # Workaround for triton bug, tensor.to doesn't unwrap constexpr values
+    DTYPE_VALUE_AS_UINT = tl.core._unwrap_if_constexpr(DTYPE_VALUE_AS_UINT)
+    bitwidth = DTYPE_VALUE_AS_UINT.primitive_bitwidth
+    uv = value.to(DTYPE_VALUE_AS_UINT, bitcast=True).to(DTYPE_PACK)
+    return flag.to(DTYPE_PACK) | (uv << bitwidth)
+
+
+@triton.jit
+def unpack_value(
+    pack,
+    DTYPE_VALUE,
+    DTYPE_VALUE_AS_UINT,
+):
+    # Workaround for triton bug, tensor.to doesn't unwrap constexpr values
+    DTYPE_VALUE = tl.core._unwrap_if_constexpr(DTYPE_VALUE)
+    DTYPE_VALUE_AS_UINT = tl.core._unwrap_if_constexpr(DTYPE_VALUE_AS_UINT)
+    bitwidth = DTYPE_VALUE_AS_UINT.primitive_bitwidth
+    value_uint = (pack >> bitwidth).to(DTYPE_VALUE_AS_UINT)
+    return value_uint.to(DTYPE_VALUE, bitcast=True)
+
+
+@triton.jit
+def unpack_flag(pack, DTYPE_FLAG):
+    return pack.to(DTYPE_FLAG)
+
+
+@triton.jit
+def exclusive_scan_decoupled_lookback(
+    scratch_base,
+    block_value,
+    index,
+    combine_fn,
+    DTYPE_VALUE_AS_UINT: tl.constexpr,
+    DTYPE_PACK: tl.constexpr,
+):
+    """Compute exclusive scan of a scalar value between blocks
+
+    Ref: https://research.nvidia.com/publication/2016-03_single-pass-parallel-prefix-scan-decoupled-look-back
+
+    scratch_base: Pointer to scratch space in global memory
+    block_value: Scalar value for this block
+    index: Scalar index of this block relative to the current scan
+    combine_fn: Function ``(value, value) -> value`` which is scanned over
+    DTYPE_VALUE_AS_UINT: A tl.uint{n} type equal in size to ``block_value``
+    DTYPE_PACK: Unsigned type twice the width of block_value
+
+    NOTE: This function is limited to values which are 32-bits or less because
+    we need to pack (value, flag) into a single unsigned int.
+    """
+    # Publish block sum so subsequent blocks don't get stuck waiting for us
+    DTYPE_VALUE = block_value.dtype
+    pack = pack_value_flag(
+        block_value,
+        tl.full(block_value.shape, 1, DTYPE_VALUE_AS_UINT),
+        DTYPE_VALUE_AS_UINT,
+        DTYPE_PACK,
+    )
+    if index > 0:
+        tl.atomic_xchg(scratch_base + index, pack, sem="relaxed")
+
+    # Calculate exclusive prefix scan
+    exclusive_prefix = tl.zeros([], DTYPE_VALUE)
+    prefix_valid = False
+    test_target = index - 1
+    while test_target >= 0:
+        # tl.atomic_load
+        flag = tl.full([], 0, DTYPE_VALUE_AS_UINT)
+        while flag == 0:
+            pack = tl.atomic_add(scratch_base + test_target, 0, sem="relaxed")
+            flag = unpack_flag(pack, DTYPE_VALUE_AS_UINT)
+
+        value = unpack_value(pack, DTYPE_VALUE, DTYPE_VALUE_AS_UINT)
+        if prefix_valid:
+            exclusive_prefix = combine_fn(value, exclusive_prefix)
+        else:
+            exclusive_prefix = value
+            prefix_valid = True
+
+        if flag == 2:
+            test_target = -1
+        else:
+            test_target = test_target - 1
+
+    # Make inclusive block sum visible to other blocks
+    if prefix_valid:
+        inclusive_prefix = combine_fn(exclusive_prefix, block_value)
+    else:
+        inclusive_prefix = block_value
+    pack = pack_value_flag(
+        inclusive_prefix,
+        tl.full([], 2, DTYPE_VALUE_AS_UINT),
+        DTYPE_VALUE_AS_UINT,
+        DTYPE_PACK,
+    )
+    tl.atomic_xchg(scratch_base + index, pack, sem="relaxed")
+    return exclusive_prefix
+
+
+@triton.jit
+def exclusive_scan_decoupled_lookback_64(scratch_base, block_value, index, combine_fn):
+    """Compute exclusive scan of a scalar value between blocks
+
+    Ref: https://research.nvidia.com/publication/2016-03_single-pass-parallel-prefix-scan-decoupled-look-back
+
+    scratch_base: Pointer to scratch space in global memory
+    block_value: Scalar value for this block, must be 64-bits wide
+    index: Scalar index of this block relative to the current scan
+    combine_fn: Function ``(value, value) -> value`` which is scanned over
+    init: Scalar value equal to the identity of combine_fn
+    """
+    # Publish block sum so subsequent blocks don't get stuck waiting for us
+    if index > 0:
+        block_value_u64 = block_value.to(tl.uint64, bitcast=True)
+        tl.store(scratch_base + 3 * index + 1, block_value_u64)
+        tl.debug_barrier()
+        flag_one = tl.full([], 1, tl.uint64)
+        tl.atomic_xchg(scratch_base + 3 * index + 0, flag_one, sem="release")
+
+    # Calculate exclusive prefix scan
+    exclusive_prefix = tl.zeros([], block_value.dtype)
+    prefix_valid = False
+    test_target = index - 1
+    while test_target >= 0:
+        flag = tl.full([], 0, tl.uint64)
+        while flag == 0:
+            flag = tl.atomic_add(scratch_base + 3 * test_target + 0, 0, sem="acquire")
+
+        value_u64 = tl.load(scratch_base + 3 * test_target + flag.to(tl.int32))
+        value = value_u64.to(block_value.dtype, bitcast=True)
+        if prefix_valid:
+            exclusive_prefix = combine_fn(value, exclusive_prefix)
+        else:
+            exclusive_prefix = value
+            prefix_valid = True
+
+        if flag == 2:
+            test_target = tl.full([], -1, index.dtype)  # Match the original type
+        else:
+            test_target = test_target - 1
+
+    # Make inclusive block sum visible to other blocks
+    if prefix_valid:
+        inclusive_prefix = combine_fn(exclusive_prefix, block_value)
+    else:
+        inclusive_prefix = block_value
+    inclusive_prefix_u64 = inclusive_prefix.to(tl.uint64, bitcast=True)
+    tl.store(scratch_base + 3 * index + 2, inclusive_prefix_u64)
+    tl.debug_barrier()
+    flag_two = tl.full([], 2, tl.uint64)
+    tl.atomic_xchg(scratch_base + 3 * index + 0, flag_two, sem="release")
+
+    return exclusive_prefix
+
+
+@triton.jit
+def frexp(x):
+    # TODO(isuruf): use inline_asm_elementwise here
+    y = libdevice.ilogb(x) + 1
+    exponent = tl.where(x == 0, 0, y)
+    mantissa = tl.where(x == 0, 0, libdevice.ldexp(x, -y))
+    return mantissa, exponent
+
+
+@triton.jit
+def _compare_and_swap_with_index(
+    x,
+    idxs,
+    rnumel,
+    flip,
+    i: tl.constexpr,
+    n_dims: tl.constexpr,
+    stable: tl.constexpr,
+    descending: tl.constexpr,
+):
+    n_outer: tl.constexpr = x.numel >> n_dims
+    shape: tl.constexpr = [n_outer * 2**i, 2, 2 ** (n_dims - i - 1)]
+
+    idtype = tl.core.get_int_dtype(bitwidth=x.dtype.primitive_bitwidth, signed=True)
+
+    y = tl.reshape(x, shape)
+    iy = y.to(idtype, bitcast=True)
+    # slice left/right with 'stride' 2**(n_dims - i - 1)
+    right_mask = tl.arange(0, 2)[None, :, None].to(idtype)
+    left_mask = (1 - right_mask).to(idtype)
+    ileft = tl.broadcast_to(tl.sum(iy * left_mask, 1).to(idtype)[:, None, :], shape)
+    iright = tl.broadcast_to(tl.sum(iy * right_mask, 1).to(idtype)[:, None, :], shape)
+    ileft = tl.reshape(ileft, x.shape)
+    iright = tl.reshape(iright, x.shape)
+    left = ileft.to(x.dtype, bitcast=True)
+    right = iright.to(x.dtype, bitcast=True)
+
+    # idx
+    y_idx = tl.reshape(idxs, shape)
+    left_idx = tl.broadcast_to(
+        tl.sum(y_idx * left_mask.to(y_idx.dtype), 1)[:, None, :], shape
+    )
+    right_idx = tl.broadcast_to(
+        tl.sum(y_idx * right_mask.to(y_idx.dtype), 1)[:, None, :], shape
+    )
+    left_idx = tl.reshape(left_idx, x.shape)
+    right_idx = tl.reshape(right_idx, x.shape)
+
+    # valid
+    if rnumel is None:
+        left_valid_mask = tl.full(x.shape, True, tl.int1)
+        right_valid_mask = tl.full(x.shape, True, tl.int1)
+    else:
+        left_valid_mask = left_idx < rnumel
+        right_valid_mask = right_idx < rnumel
+
+    # actual compare-and-swap
+    ix = x.to(idtype, bitcast=True)
+
+    # sort treats nan as having the higher value. comparisons with nan always return False.
+    # to align with sort semantics, we need to update descending to check if right_isnan,
+    # and ascending to check if left_isnan.
+    left_isnan = left != left
+    right_isnan = right != right
+
+    if descending:
+        cond = left < right
+        if is_floating(left):
+            if not stable:
+                cond = cond | right_isnan
+            else:
+                cond = cond | (right_isnan & (~left_isnan))
+
+    else:
+        cond = left > right
+        if is_floating(left):
+            if not stable:
+                cond = cond | left_isnan
+            else:
+                cond = cond | (left_isnan & (~right_isnan))
+
+    if stable:
+        # When stable sorting, tie break by index
+        eq = left == right
+        if is_floating(left):
+            eq = eq | (left_isnan & right_isnan)
+        cond = cond | (eq & (left_idx > right_idx))
+
+    cond = (right_valid_mask > left_valid_mask) | (
+        (right_valid_mask == left_valid_mask) & cond
+    )
+    cond = (cond ^ flip).to(tl.int1)
+    ret = ix ^ tl.where(cond, ileft ^ iright, tl.zeros_like(ix))
+    new_idxs = idxs ^ tl.where(cond, left_idx ^ right_idx, tl.zeros_like(idxs))
+
+    return ret.to(x.dtype, bitcast=True), new_idxs
+
+
+@triton.jit
+def _bitonic_merge_with_index(
+    x,
+    idxs,
+    rnumel,
+    stage: tl.constexpr,
+    alternating: tl.constexpr,
+    n_dims: tl.constexpr,
+    stable: tl.constexpr,
+    descending: tl.constexpr,
+):
+    n_outer: tl.constexpr = x.numel >> n_dims
+    tl.static_assert(stage <= n_dims)
+    # flip denotes whether to re-arrange sub-sequences of elements in ascending or
+    # descending order.
+    # if flip = 00000000... then all elements will be re-arranged ascendingly at this stage
+    # if flip = 00110011... then all the elements will be re-arranged alternatingly (with
+    # a stride of 2) at this stage
+    if alternating:
+        shape: tl.constexpr = [n_outer * 2 ** (n_dims - 1 - stage), 2, 2**stage]
+        flip = tl.reshape(
+            tl.broadcast_to(tl.arange(0, 2)[None, :, None], shape), x.shape
+        )
+    else:
+        flip = False
+    # perform `stage` rounds of `compare-and-swap`
+    for i in tl.static_range(stage):
+        x, idxs = _compare_and_swap_with_index(
+            x, idxs, rnumel, flip, i + (n_dims - stage), n_dims, stable, descending
+        )
+    return x, idxs
+
+
+@triton.jit
+def sort_with_index(
+    x,  # value
+    idxs,  # index
+    rnumel,  # number of elements
+    dim: tl.constexpr = None,
+    stable: tl.constexpr = tl.constexpr(False),
+    descending: tl.constexpr = tl.constexpr(False),
+):
+    x, idxs = tl.broadcast(x, idxs)
+    # handle default dimension or check that it is the most minor dim
+    _dim: tl.constexpr = len(x.shape) - 1 if dim is None else dim
+    tl.static_assert(
+        _dim == len(x.shape) - 1, "only minor dimension is currently supported"
+    )
+    # iteratively run bitonic merge-sort steps
+    n_dims: tl.constexpr = _log2(x.shape[_dim])
+
+    for i in tl.static_range(1, n_dims + 1):
+        x, idxs = _bitonic_merge_with_index(
+            x,
+            idxs,
+            rnumel,
+            i,
+            alternating=i < n_dims,
+            n_dims=n_dims,
+            stable=stable,
+            descending=descending,
+        )
+    return x, idxs
+
+
+@triton.jit
+def select_one(x, mask, dim, keep_dims=False):
+    idtype = tl.core.get_int_dtype(x.dtype.primitive_bitwidth, signed=False)
+    ix = x.to(idtype, bitcast=True)
+    iy = tl.sum(ix * mask, dim, keep_dims=keep_dims)
+    return iy.to(idtype).to(x.dtype, bitcast=True)
+
+
+@triton.jit
+def x_grid_barrier(sem):
+    """
+    Wait for all other thread blocks in grid sharing same y/z program_id
+    to reach this barrier before returning.
+
+    Args:
+        sem: an uint32 semaphores, zero or 0x80000000 initialized.  Must be unique to each y/z program ID.
+    """
+    # ensure stores before this are visible
+    tl.debug_barrier()
+
+    one_i32 = 1
+    one_u32 = one_i32.to(tl.uint32)  # type: ignore[attr-defined]
+    expected = tl.num_programs(0).to(tl.uint32)
+    if tl.program_id(0) == 0:
+        nb = 0x80000000 - (expected - one_u32)
+    else:
+        nb = one_u32
+
+    old_arrive = tl.atomic_add(sem, nb, sem="release")
+
+    bar_flipped = False
+    while not bar_flipped:
+        # want a `ld.acquire.gpu.u32 $0,[$1];` but Triton doesn't have it
+        current_arrive = tl.atomic_add(sem, 0, sem="acquire")
+        # current_arrive = tl.load(sem, volatile=True)
+        bar_flipped = ((old_arrive ^ current_arrive) & 0x80000000) != 0
+
+    # TODO(jansel): is this needed?
+    tl.debug_barrier()
+
+
+def triton_builtin(f: Callable[..., _T]) -> Callable[..., _T]:
+    """
+    Decorator to mark a function as a Triton built-in function.  These functions
+    are evaluated at compile time.
+
+    Args:
+        f (function): The function to be marked as a Triton built-in.
+
+    Returns:
+        function: The same function, marked as a Triton built-in.
+    """
+    if builtins_use_semantic_kwarg:
+        # support Triton before and after https://github.com/triton-lang/triton/pull/7054
+        # and after https://github.com/triton-lang/triton/pull/7239
+        def wrapper(*args, _semantic, **kwargs):
+            kwargs["_builder"] = _semantic
+            return f(*args, **kwargs)
+    else:
+        wrapper = f  # type: ignore[assignment]
+
+    wrapper.__triton_builtin__ = True  # type: ignore[attr-defined]
+    return wrapper
+
+
+@triton_builtin
+def constexpr_next_power_of_2(
+    n: tl.constexpr, *, _builder: object = None
+) -> tl.constexpr:
+    """
+    A version triton.next_power_of_two that can be used within a kernel on constants.
+    """
+    assert isinstance(n, tl.constexpr)
+    return tl.constexpr(triton.next_power_of_2(n.value))
+
+
+@triton_builtin
+def if_mask(mask: Any, val, *, _builder: object = None) -> tl.constexpr:
+    """
+    Work around triton compile error: `ValueError: `other` cannot be provided without `mask``
+    A compile-time to check to return either `val` or `None` depending on the value of mask.
+    """
+    if isinstance(mask, tl.constexpr) and mask.value is None:
+        return tl.constexpr(None)
+    return val
+
+
+@triton.jit
+def inline_asm_pack(x, pack: tl.constexpr):
+    """Ravel to 1D and pad (via join with zeros) so numel is divisible by pack."""
+    result = x.ravel()
+    # Only pad when the block size is smaller than pack. When block >= pack
+    # the numel is already divisible by pack (both are powers of 2).
+    n_pad: tl.constexpr = _log2(pack) - _log2(result.numel)
+    for _ in tl.static_range(n_pad):
+        result = tl.reshape(
+            tl.join(result, tl.zeros_like(result)), (result.shape[0] * 2,)
+        )
+    return result
+
+
+@triton.jit
+def inline_asm_unpack(x, orig, pack: tl.constexpr):
+    """Unpad and reshape back to orig's shape."""
+    result = x
+    n_pad: tl.constexpr = _log2(pack) - _log2(orig.numel)
+    for _ in tl.static_range(n_pad):
+        result, _ = tl.split(tl.reshape(result, (result.shape[0] // 2, 2)))
+    return tl.reshape(result, orig.shape)

outputs/audit_venv/lib/python3.11/site-packages/torch/_inductor/runtime/triton_lazy_compile.py

@@ -0,0 +1,243 @@
+"""
+Lazy Triton kernel compilation for C++ wrapper.
+
+This module provides functionality to compile and autotune Triton kernels at runtime
+when the C++ wrapper is used without autotune_at_compile_time.
+
+The workflow is:
+1. At model initialization: Call start_kernel_compile() for all kernels to start
+   parallel compilation using multi-process async_compile
+2. At kernel execution time: Call run_triton_kernel_with_autotune() which waits
+   for the specific kernel to be ready, then runs it with autotuning
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import logging
+import re
+from typing import Any
+
+from .triton_heuristics import CachingAutotuner
+
+
+log = logging.getLogger(__name__)
+
+
+@dataclasses.dataclass
+class TritonKernelCompileResult:
+    cubin_path: str
+    mangled_name: str
+    num_warps: int
+    shared_mem: int
+    xblock: int
+    yblock: int
+    zblock: int
+    r0block: int
+    rsplit: int
+    rsplit_size: int
+    config_index: int | None
+    global_scratch: int | None
+    profile_scratch: int | None
+
+
+_async_compile: Any = None
+
+
+def _get_async_compile() -> Any:
+    """Get or create the shared AsyncCompile instance."""
+    global _async_compile
+    if _async_compile is None:
+        from torch._inductor.async_compile import AsyncCompile
+
+        _async_compile = AsyncCompile()
+    return _async_compile
+
+
+def _wrap_tma_args(args: list[Any], kernel_fn: CachingAutotuner) -> list[Any]:
+    """Wrap tensor args with TMA descriptors where the signature requires them."""
+    signature = kernel_fn.triton_meta.get("signature", {})
+    sig_items = list(signature.items())
+
+    # Track args index separately from sig_items index since the signature
+    # may include constexpr entries that are not present in args.
+    tma_indices = []
+    arg_idx = 0
+    for name, sig_type in sig_items:
+        if isinstance(sig_type, str) and sig_type == "constexpr":
+            continue
+        if isinstance(sig_type, str) and sig_type == "nvTmaDesc":
+            raise RuntimeError(
+                f"nvTmaDesc (experimental TMA API) is not supported in lazy compile "
+                f"for arg '{name}'. Use the stable tensordesc API instead."
+            )
+        if isinstance(sig_type, str) and sig_type.startswith("tensordesc<"):
+            tma_indices.append((arg_idx, name, sig_type))
+        arg_idx += 1
+
+    if not tma_indices:
+        return args
+
+    from triton.tools.tensor_descriptor import TensorDescriptor
+
+    wrapped = list(args)
+    for arg_idx, name, sig_type in tma_indices:
+        if arg_idx >= len(wrapped):
+            raise RuntimeError(
+                f"TMA arg index {arg_idx} for '{name}' exceeds arg count {len(wrapped)}"
+            )
+        tensor = wrapped[arg_idx]
+        # Parse block_shape from tensordesc<dtype[dim0, dim1, ...]>
+        match = re.match(r"tensordesc<[^[]*\[([^\]]*)\]", sig_type)
+        if match:
+            block_shape = [int(x.strip()) for x in match.group(1).split(",")]
+            wrapped[arg_idx] = TensorDescriptor.from_tensor(tensor, block_shape)
+
+    return wrapped
+
+
+def start_kernel_compile(
+    pending_kernels: dict[str, Any], kernel_name: str, kernel_source: str
+) -> None:
+    """
+    This function is called from C++ at model initialization time for each kernel.
+    It starts the compilation in a background process but does NOT wait for it.
+    The actual kernel execution happens later in run_triton_kernel_with_autotune().
+
+    The pending_kernels dict is per-module, created in C++ and passed through
+    to avoid global state collisions across compiled modules.
+    """
+    if kernel_name in pending_kernels:
+        return
+
+    async_compile = _get_async_compile()  # noqa: F841 (used by eval below)
+
+    # Evaluate the kernel source to get the Future or CachingAutotuner
+    # The kernel_source is like: async_compile.triton('name', '''...''', ...)
+    kernel_obj = eval(kernel_source.strip())  # noqa: S307
+
+    pending_kernels[kernel_name] = kernel_obj
+
+
+def run_triton_kernel_with_autotune(
+    pending_kernels: dict[str, Any],
+    kernel_name: str,
+    stream: Any,
+    args: list[Any],
+) -> TritonKernelCompileResult:
+    """
+    Run a Triton kernel with full autotuning using actual tensor arguments.
+    """
+    from torch._inductor.codecache import CodeCacheFuture, CudaKernelParamCache
+    from torch._inductor.runtime.triton_heuristics import config_to_dict
+
+    if kernel_name not in pending_kernels:
+        raise RuntimeError(f"Kernel {kernel_name} not found in pending kernels.")
+    kernel_obj = pending_kernels[kernel_name]
+
+    if isinstance(kernel_obj, CodeCacheFuture):
+        kernel_fn = kernel_obj.result()
+    elif isinstance(kernel_obj, CachingAutotuner):
+        kernel_fn = kernel_obj
+    else:
+        raise RuntimeError(f"Unexpected kernel object type: {type(kernel_obj)}")
+
+    assert isinstance(kernel_fn, CachingAutotuner)
+
+    inductor_meta = kernel_fn.inductor_meta
+    inductor_meta["store_cubin"] = True
+
+    # For TMA kernels, wrap tensor args with TMA descriptors
+    args = _wrap_tma_args(args, kernel_fn)
+
+    # Run the kernel with the provided arguments
+    # This will trigger autotuning if there are multiple configs
+    kernel_fn.run(*args, stream=stream)
+    if not kernel_fn.launchers:
+        raise RuntimeError("Kernel run did not produce any launchers")
+    launcher = kernel_fn.launchers[0]
+
+    cached_params: dict[str, Any] | None = CudaKernelParamCache.get(kernel_name)
+    if cached_params is None:
+        raise RuntimeError(f"Failed to get cached params for kernel {kernel_name}")
+
+    from torch._inductor.codecache import get_cpp_wrapper_cubin_path_name
+
+    cubin_path_name = get_cpp_wrapper_cubin_path_name()
+    for key_name in (cubin_path_name, "mangled_name", "num_warps", "shared_mem"):
+        if key_name not in cached_params:
+            raise RuntimeError(
+                f"{key_name} not found in cached params for {kernel_name}"
+            )
+    cubin_path = cached_params[cubin_path_name]
+    mangled_name = cached_params["mangled_name"]
+    num_warps = cached_params["num_warps"]
+    shared_mem = cached_params["shared_mem"]
+
+    config = config_to_dict(launcher.config) if launcher.config else {}
+
+    # For combo/foreach kernels, the autotuned config may have empty kwargs
+    # (e.g., the foreach heuristic only tunes num_warps, not XBLOCK).
+    # In that case, use the default_config from combo_grid_meta
+    combo_grid_meta = inductor_meta.get("combo_grid_meta") if inductor_meta else None
+    default_config = combo_grid_meta.get("default_config") if combo_grid_meta else None
+    if default_config:
+        config = {**default_config, **config}
+
+    xblock = config.get("XBLOCK", 128)
+    yblock = config.get("YBLOCK", 1)
+    zblock = config.get("ZBLOCK", 1)
+    r0block = config.get("R0_BLOCK", 1)
+    rsplit = config.get("RSPLIT", 1)
+    rsplit_size = config.get("RSPLIT_SIZE", 1)
+
+    config_index = None
+    grid_type = inductor_meta.get("grid_type") if inductor_meta else None
+    if grid_type == "PrecomputedGrid" and inductor_meta:
+        # PrecomputedGrid selects one of precomputed_grids. We use config_index
+        # to remember which grid is chosen.
+        precomputed_grids = inductor_meta.get("precomputed_grids", [])
+        for idx, entry in enumerate(precomputed_grids):
+            entry_config = entry.get("config", {})
+            if all(config.get(k) == v for k, v in entry_config.items()):
+                config_index = idx
+                break
+
+    global_scratch: int | None = cached_params.get("global_scratch")
+    profile_scratch: int | None = cached_params.get("profile_scratch")
+
+    log.debug(
+        "Successfully autotuned Triton kernel: cubin_path=%s, mangled_name=%s, "
+        "num_warps=%d, shared_mem=%d, xblock=%d, yblock=%d, zblock=%d, r0block=%d, "
+        "rsplit=%d, rsplit_size=%d, config_index=%s, global_scratch=%s, profile_scratch=%s",
+        cubin_path,
+        mangled_name,
+        num_warps,
+        shared_mem,
+        xblock,
+        yblock,
+        zblock,
+        r0block,
+        rsplit,
+        rsplit_size,
+        config_index,
+        global_scratch,
+        profile_scratch,
+    )
+
+    result = TritonKernelCompileResult(
+        cubin_path=cubin_path,
+        mangled_name=mangled_name,
+        num_warps=num_warps,
+        shared_mem=shared_mem,
+        xblock=xblock,
+        yblock=yblock,
+        zblock=zblock,
+        r0block=r0block,
+        rsplit=rsplit,
+        rsplit_size=rsplit_size,
+        config_index=config_index,
+        global_scratch=global_scratch,
+        profile_scratch=profile_scratch,
+    )
+    return result

outputs/audit_venv/lib/python3.11/site-packages/torch/_inductor/script.ld

@@ -0,0 +1,8 @@
+SECTIONS {
+  /* By default, in LLD 16, .lrodata is placed immediately after .rodata.
+   * However, .lrodata can be very large in our compiled models, which leads to
+   * relocation out-of-range errors for relative relocations. So we place it
+   * after other the sections that are referenced from .text using relative
+   * relocations. This is the default behavior in GNU ld. */
+  .lrodata : { *(.lrodata) }
+ } INSERT AFTER .bss;

outputs/audit_venv/lib/python3.11/site-packages/torch/_inductor/shape_propagation.py

@@ -0,0 +1,148 @@
+import functools
+from collections.abc import Callable, Sequence
+from typing import Protocol
+
+import sympy
+
+import torch
+
+from .virtualized import OpsValue, V
+
+
+BlockShapeType = Sequence[int | str] | None
+
+
+class ShapeVar(Protocol):
+    @property
+    def shape(self) -> BlockShapeType: ...
+
+
+ShapeArg = ShapeVar | torch.types.Number | str | OpsValue | torch.dtype
+
+# Inputs need to be cacheable (e.g., not a CSEVar) in order for the cache to be effective
+# So first decompose CSEVars -> tuple before calling this
+
+
+@functools.lru_cache(None)
+def get_broadcasted_shape(a: BlockShapeType, b: BlockShapeType) -> BlockShapeType:
+    assert isinstance(a, Sequence)
+    assert isinstance(b, Sequence)
+    if len(a) > len(b):
+        return get_broadcasted_shape(a, (*[1] * (len(a) - len(b)), *b))
+    elif len(a) < len(b):
+        b, a = a, b
+        return get_broadcasted_shape(a, (*[1] * (len(a) - len(b)), *b))
+    else:
+
+        def _get_broadcasted_dim(d1: int | str, d2: int | str) -> int | str:
+            if str(d1) == "1":
+                return d2
+            elif str(d2) == "1":
+                return d1
+            assert str(d1) == str(d2)
+            return d1
+
+        return tuple(_get_broadcasted_dim(d1, d2) for d1, d2 in zip(a, b))
+
+
+def broadcast_shapes_for_args(args: Sequence[ShapeArg]) -> BlockShapeType:
+    result_shape: BlockShapeType = None
+
+    for arg in args:
+        if hasattr(arg, "shape"):
+            shape = arg.shape
+            if shape is None:
+                return None
+            elif result_shape is None:
+                result_shape = tuple(shape)
+            else:
+                result_shape = get_broadcasted_shape(result_shape, tuple(shape))
+        elif isinstance(arg, (int, float)):
+            if result_shape is None:
+                result_shape = ()
+        elif isinstance(arg, torch.dtype):
+            continue
+        else:
+            from torch._inductor.loop_body import LoopBody, LoopBodyBlock
+
+            if isinstance(arg, (LoopBodyBlock, LoopBody, OpsValue)):
+                # TODO: fix me
+                return None
+            raise TypeError(f"Unknown type: {type(arg)}")
+
+    return result_shape
+
+
+class ShapePropagationOpsHandler:
+    """
+    Propagate shape from args to output
+    """
+
+    @staticmethod
+    def constant(value: torch.types.Number, dtype: torch.dtype) -> BlockShapeType:
+        # TritonKernelOverrides.constant uses tl.full with shape=[1]*ndim for all types
+        from torch._inductor.codegen.triton import TritonKernel
+
+        if isinstance(V.kernel, TritonKernel):
+            ndim = V.kernel.triton_tensor_ndim()
+            return tuple([1] * ndim)
+        else:
+            return ()
+
+    @staticmethod
+    def store_reduction(name: str, index: int, value: ShapeArg) -> None:
+        return None
+
+    @staticmethod
+    def reduction(
+        dtype: torch.dtype,
+        src_dtype: torch.dtype,
+        reduction_type: str,
+        value: ShapeArg | tuple[ShapeArg, ...],
+    ) -> BlockShapeType | tuple[BlockShapeType, ...]:
+        raise NotImplementedError
+
+    @staticmethod
+    def store(name: str, index: int, value: ShapeArg, mode: str | None = None) -> None:
+        return None
+
+    @staticmethod
+    def to_dtype(
+        value: ShapeVar,
+        dtype: torch.dtype,
+        src_dtype: torch.dtype | None = None,
+        use_compute_types: bool = True,
+    ) -> BlockShapeType:
+        return value.shape
+
+    @staticmethod
+    def dot(a: sympy.Expr, b: sympy.Expr) -> BlockShapeType:
+        from torch._inductor.codegen.triton import TritonKernel
+
+        assert isinstance(V.kernel, TritonKernel), "dot supports Triton only"
+        return ("YBLOCK", "XBLOCK")
+
+    @staticmethod
+    def index_expr(expr: sympy.Expr, dtype: torch.dtype) -> BlockShapeType:
+        # shape is implicitly embedded in expr.
+        return None
+
+    @staticmethod
+    def load_seed(name: str, offset: int) -> BlockShapeType:
+        return ()
+
+    @staticmethod
+    def indirect_indexing(
+        var: ShapeArg,
+        size: sympy.Expr | int,
+        check: bool = True,
+        wrap_neg: bool = True,
+    ) -> None:
+        return None
+
+    def __getattr__(self, name: str) -> Callable[..., BlockShapeType]:
+        return lambda *args, **kwargs: broadcast_shapes_for_args(args)
+
+    @staticmethod
+    def device_assert_async(cond: ShapeArg, msg: str) -> None:
+        return None

outputs/audit_venv/lib/python3.11/site-packages/torch/_inductor/sizevars.py

@@ -0,0 +1,1193 @@
+# mypy: allow-untyped-defs
+import functools
+import itertools
+import logging
+from collections.abc import Callable, Iterable, Sequence
+from typing import Any, cast
+
+import sympy
+from sympy import Expr
+
+from torch import SymInt
+from torch.fx.experimental._size_hinting import (
+    _guarding_hint_or_throw_base,
+    _maybe_realize_expr,
+    _optimization_hint_base,
+)
+from torch.fx.experimental.symbolic_shapes import (
+    free_symbols,
+    free_unbacked_symbols,
+    IterateExprs,
+    ShapeEnv,
+    SymNode,
+)
+from torch.utils._ordered_set import OrderedSet
+from torch.utils._sympy.functions import FloorDiv, Mod, ModularIndexing
+from torch.utils._sympy.symbol import symbol_is_type, SymT
+from torch.utils._sympy.value_ranges import IntInfinity, ValueRanges
+
+from . import config
+from .runtime.runtime_utils import is_power_of_2
+from .utils import (
+    has_free_symbols,
+    sympy_index_symbol,
+    sympy_index_symbol_with_prefix,
+    sympy_subs,
+    VarRanges,
+)
+from .virtualized import V
+
+
+log = logging.getLogger(__name__)
+
+# Symbols created by CppTemplateKernel.slice_nd → parse_expr_with_index_symbols
+# and are internal to C++ GEMM codegen (not tracked by ShapeEnv).
+_GEMM_TEMPLATE_SYMBOL_NAMES = OrderedSet(
+    [
+        "m_start",
+        "m_end",
+        "n_start",
+        "n_end",
+        "k_start",
+        "k_end",
+        "m_size",
+        "n_size",
+        "m_offset",
+        "m_start_unsliced",
+        "m_end_unsliced",
+        "m_size_unsliced",
+    ]
+)
+
+
+def statically_known_true(
+    shape_env: ShapeEnv,
+    expr: sympy.Basic | bool,
+    axioms: tuple[sympy.Expr] | None = None,
+    var_to_range: tuple[tuple[sympy.Symbol, ValueRanges[Any]]] | None = None,
+) -> bool:
+    if expr in (True, False):
+        return bool(expr)
+
+    try:
+        simplified = shape_env._maybe_evaluate_static(
+            expr,
+            axioms=axioms,
+            var_to_range=var_to_range,
+        )
+        if simplified is not None:
+            return bool(simplified)
+    except Exception:
+        log.debug("Could not simplify  %s", expr, exc_info=True)
+
+    return False
+
+
+# This class is a little awkward, because ShapeEnv is doing most of the heavy
+# lifting and in some cases we should be directly passing through to ShapeEnv,
+# but there is some extra inductor logic that needs to be handled here
+class SizeVarAllocator:
+    """
+    A class that manages symbolic size variables and their relationships.
+
+    This class works with the ShapeEnv to handle symbolic shape expressions,
+    simplify them, and provide utilities for guarding, checking, and evaluating
+    symbolic expressions. It also manages precomputed replacements and stride
+    calculations for tensor operations.
+    """
+
+    def __init__(self, shape_env=None) -> None:
+        super().__init__()
+        # Note: this can lead to bugs. Reasoning APIs depends on existing information in
+        # in the shape_env. For example! var_to_ranges can't be empty!
+        if shape_env is None:
+            shape_env = ShapeEnv()
+        self.shape_env = shape_env
+        self.backed_var_to_val = self.shape_env.backed_var_to_val
+        self.var_to_hint_override = self.shape_env.var_to_hint_override
+        self.replacements: dict[sympy.Symbol, Expr] = self.shape_env.replacements
+        self.unbacked_replacements: dict[Expr, Expr] | None = None
+        # Maps of dynamic sizes that have to be precomputed on the host to the kernel args.
+        # The basic idea is if we have some complicated sympy expression
+        # f(s0), we may choose to precompute it on the host and then replace
+        # all occurrences of that sympy expression with ps0, so that when we
+        # codegen we simply reference ps0 directly without repeating
+        # f(s0).  Unlike regular size variables, ps variables cannot be
+        # guarded upon; so if we are asked to guard on a Sympy expression
+        # which potentially could have already had a precomputed replacement
+        # on it, we are obligated to invert the precomputed replacements
+        # (inv_precomputed_replacements).
+        self.precomputed_replacements: dict[Expr, sympy.Symbol] = {}
+        self.inv_precomputed_replacements: dict[sympy.Symbol, Expr] = {}
+        self.stride_vars = self.make_stride_vars_cache()
+        self.simplify_with_ranges = self.make_simplify_with_ranges_cache()
+        self._simplify_loops = self.make_simplify_loops_cache()
+
+    def simplify(self, expr: Expr):
+        return sympy.expand(expr).xreplace(self.replacements)
+
+    def make_simplify_with_ranges_cache(self) -> Callable[[Expr, VarRanges], Expr]:
+        """
+        self._simplify_with_ranges() can be expensive, cache its results
+        """
+        cache: dict[tuple[Any, ...], Expr] = {}
+        replacement_count = len(self.replacements)
+
+        def simplify_with_ranges(expr: Expr, var_ranges: VarRanges) -> Expr:
+            nonlocal replacement_count
+            if replacement_count != len(self.replacements):
+                # new replacements invalidates cached results
+                cache.clear()
+                replacement_count = len(self.replacements)
+            key = (expr, *var_ranges.items())
+            result = cache.get(key)
+            if result is None:
+                result = self._simplify_with_ranges(expr, var_ranges)
+                cache[key] = result
+                if result != expr:
+                    cache[(result, *var_ranges.items())] = result
+            return result
+
+        return simplify_with_ranges
+
+    def make_simplify_loops_cache(self):
+        """
+        self._simplify_with_ranges() can be expensive, cache its results
+        """
+        cache: dict[tuple[Any, ...], Any] = {}
+        replacement_count = len(self.replacements)
+
+        def simplify_loops(index_vars, sizes, index_formulas):
+            nonlocal replacement_count
+            if replacement_count != len(self.replacements):
+                # new replacements invalidates cached results
+                cache.clear()
+                replacement_count = len(self.replacements)
+            key = (*index_vars, *sizes, *index_formulas)
+            result = cache.get(key)
+            if result is None:
+                result = self._simplify_loops_impl(index_vars, sizes, index_formulas)
+                cache[key] = result
+            return result
+
+        return simplify_loops
+
+    def _simplify_with_ranges(self, expr: Expr, var_ranges: VarRanges) -> Expr:
+        """
+        Simplify indexing expression with knowledge of the ranges of
+        iteration variables.
+        """
+
+        expr = join_dimensions(self.simplify(expr))
+        original_expr = expr
+
+        var_to_range = dict(self.shape_env.var_to_range)
+        var_to_range.update(
+            {
+                k: ValueRanges(
+                    0, max(0, v - 1) if not has_free_symbols([v]) else IntInfinity()
+                )
+                for k, v in var_ranges.items()
+            }
+        )
+        for var in expr.free_symbols:
+            if var not in var_to_range:
+                var_to_range[var] = ValueRanges(0, IntInfinity())
+
+        var_to_range_tuple = cast(
+            tuple[tuple[sympy.Symbol, ValueRanges[sympy.Expr]]],
+            tuple(var_to_range.items()),
+        )
+
+        axioms = []
+        for var, upper_bound in var_ranges.items():
+            axioms.append(0 <= var)
+            axioms.append(var < upper_bound)
+        axioms = tuple(axioms) + self.shape_env.get_axioms()
+
+        def statically_known(expr):
+            evaluated = self.shape_env._maybe_evaluate_static(
+                expr,
+                # pyrefly: ignore [bad-argument-type]
+                axioms=axioms,
+                var_to_range=var_to_range_tuple,
+            )
+            return bool(evaluated)
+
+        def remove_zero_terms(base, divisor):
+            """Symbols smaller than the divisor are zero"""
+            if not statically_known(base >= 0):
+                return base
+
+            for v in base.free_symbols:
+                if v in var_ranges:
+                    rest = sympy.Wild("_rest", exclude=[v])
+                    m = base.match(v + rest)
+                    if m and v not in m[rest].free_symbols:
+                        # v can be removed if it doesn't affect the FloorDiv.
+                        # rest is always a multiple of gcd(rest, divisor), so
+                        # rest % divisor is also a multiple of that gcd. The
+                        # worst case is rest % divisor == divisor - gcd, so
+                        # adding v is safe when v < gcd.
+                        gcd = sympy.gcd(m[rest], divisor)
+                        if statically_known(v < gcd):
+                            base = m[rest]
+            return base
+
+        def visit_indexing_div(base, divisor):
+            base = remove_zero_terms(base, divisor)
+            if statically_known(base >= 0) and statically_known(base < divisor):
+                return sympy.S.Zero
+            # FloorDiv(ModularIndexing(b, d1, m), d2) = ModularIndexing(b, d1*d2, m//d2)
+            if isinstance(base, ModularIndexing) and isinstance(divisor, sympy.Integer):
+                b, d1, m = base.args
+                if m % divisor == 0:
+                    return ModularIndexing(b, d1 * divisor, FloorDiv(m, divisor))
+            return FloorDiv(base, divisor)
+
+        def visit_modular_indexing(base, divisor, modulus):
+            base = remove_zero_terms(base, divisor)
+
+            can_remove_mod = statically_known(base >= 0) and statically_known(
+                base < modulus * divisor
+            )
+
+            if can_remove_mod:
+                return FloorDiv(base, divisor)
+            return ModularIndexing(base, divisor, modulus)
+
+        if expr.has(ModularIndexing):
+            expr = expr.replace(
+                ModularIndexing(
+                    sympy.Wild("base", integer=True),
+                    sympy.Wild("divisor", integer=True),
+                    sympy.Wild("modulus", integer=True),
+                ),
+                visit_modular_indexing,
+            )
+
+        if expr.has(FloorDiv):
+            expr = expr.replace(
+                FloorDiv(
+                    sympy.Wild("base", integer=True),
+                    sympy.Wild("divisor", integer=True),
+                ),
+                visit_indexing_div,
+            )
+
+        if expr != original_expr:
+            return self._simplify_with_ranges(expr, var_ranges)
+        return expr
+
+    def _simplify_loops_impl(
+        self, index_vars: list[sympy.Symbol], sizes, index_formulas
+    ):
+        """
+        Try to remove as many axis from loop iterations as possible, by:
+            1) removing size==1 dimensions
+            2) fuse contiguous dimensions into a single loop
+            If channel_last = True, we will prevent the last dim fused with other dims
+        """
+        sizes = list(map(self.simplify, sizes))
+
+        strides = [
+            # index_formulas may contain boolean expressions (e.g. s0 < 10),
+            # for which "strides" don't make sense so we ignore them here.
+            # NOTE: These expressions may still block merging dims in the sound
+            # substitution test performed in can_merge_dims.
+            (
+                self.stride_vars(x, index_vars)
+                if isinstance(x, sympy.Expr)
+                else [0] * len(index_vars)
+            )
+            for x in index_formulas
+        ]
+        assert len(sizes) == len(strides[0]), (len(sizes), len(strides[0]))
+
+        for i in range(len(sizes)):
+            if sizes[i] == 1:
+                # remove dim
+                sizes[i] = None
+
+        def can_merge_dims(a, b):
+            for k in range(len(strides)):
+                if self.simplify(strides[k][a] * sizes[a]) == self.simplify(
+                    strides[k][b]
+                ):
+                    # approximate test passed, try sound version
+                    va = index_vars[a]
+                    vb = index_vars[b]
+                    m1 = sympy_index_symbol("_merge_tester1")
+                    m2 = sympy_index_symbol("_merge_tester2")
+                    # NOTE: can't sub vb=0 here in case va * vb appears in the expression,
+                    # in which case both expr1 and expr2 would be zero!
+                    expr1 = sympy_subs(index_formulas[k], {va: m1 * sizes[a], vb: m2})
+                    expr2 = sympy_subs(index_formulas[k], {va: 0, vb: (m1 + m2)})
+                    if self.simplify(expr1) == self.simplify(expr2):
+                        continue
+                return False
+            return True
+
+        changed = True
+        while changed:
+            changed = False
+            for i, j in itertools.product(
+                reversed(range(len(sizes))), reversed(range(len(sizes)))
+            ):
+                if i == j or sizes[i] is None or sizes[j] is None:
+                    continue
+                if can_merge_dims(i, j):
+                    changed = True
+                    sizes[i] = sizes[i] * sizes[j]
+                    sizes[j] = None
+
+        def reindex(index):
+            it = list(reversed(index))
+            new_index = []
+            for size in sizes:
+                if size is None:
+                    new_index.append(sympy.S.Zero)
+                else:
+                    new_index.append(it.pop())
+            assert not it
+            return new_index
+
+        def prune(index):
+            assert len(index) == len(sizes)
+            return [i for i, s in zip(index, sizes) if s is not None]
+
+        return [x for x in sizes if x is not None], reindex, prune
+
+    # Note - [On Statically Known]
+    # The statically_known_* family of functions below NEVER guard, they could return True if the
+    # asked questions can be answered without guarding otherwise they return False.
+    # Those are similar to statically_known_true in symbolic_shapes.py but operate on sympy
+    # expressions instead of symnodes.
+    def statically_known_true(self, expr: sympy.Basic | bool) -> bool:
+        """
+        Returns true if an expression is always true (symbolically or via guards),
+        false otherwise. Never add guards, or throw data dependent errors.
+        """
+        return statically_known_true(self.shape_env, expr)
+
+    def statically_known_equals(self, left: Expr | int, right: Expr | int) -> bool:
+        """
+        Returns a bool indicating if it is sound to optimize as if left and right are equal.
+        """
+        return self.statically_known_true(sympy.Eq(left, right))  # type: ignore[arg-type]
+
+    def statically_known_list_equals(
+        self, left: Sequence[Expr], right: Sequence[Expr]
+    ) -> bool:
+        """
+        Returns a bool indicating if it is sound to optimize as if left and right lists are equal.
+        """
+        return len(left) == len(right) and all(
+            self.statically_known_equals(l, r) for l, r in zip(left, right)
+        )
+
+    def statically_known_leq(self, left: Expr, right: Expr | int) -> bool:
+        """
+        Returns a bool indicating if it is sound to optimize as if left is less than or equal to right.
+        """
+        expr = left <= right
+        return self.statically_known_true(expr)
+
+    def statically_known_geq(self, left: Expr, right: Expr | int) -> bool:
+        """
+        Returns a bool indicating if it is sound to optimize as if left is greater than or equal to right.
+        """
+        expr = left >= right
+        return self.statically_known_true(expr)
+
+    def statically_known_lt(self, left: Expr, right: Expr | int) -> bool:
+        """
+        Returns a bool indicating if it is sound to optimize as if left is less than right.
+        """
+        expr = left < right
+        return self.statically_known_true(expr)
+
+    def statically_known_gt(self, left: Expr, right: Expr | int) -> bool:
+        """
+        Returns a bool indicating if it is sound to optimize as if left is greater than right.
+        """
+        expr = left > right
+        return self.statically_known_true(expr)
+
+    def _is_multiple_of(self, numerator: Expr, denominator: int) -> bool:
+        """
+        Structural divisibility check: returns True only if numerator is
+        provably a multiple of denominator.  Recurses over sympy expression
+        structure before falling back to statically_known_true.
+        """
+        # Rule 1 — concrete value
+        if isinstance(numerator, (int, sympy.Integer)):
+            return int(numerator) % denominator == 0
+
+        # Rule 2 — product: any factor divisible → product divisible
+        if isinstance(numerator, sympy.Mul):
+            for factor in numerator.args:
+                if self._is_multiple_of(factor, denominator):
+                    return True
+            # Also check if combined constant factors are divisible
+            const = 1
+            for factor in numerator.args:
+                if isinstance(factor, (int, sympy.Integer)):
+                    const *= int(factor)
+            if const != 1 and const % denominator == 0:
+                return True
+
+        # Rule 3 — sum: all terms divisible → sum divisible
+        if isinstance(numerator, sympy.Add):
+            if all(self._is_multiple_of(term, denominator) for term in numerator.args):
+                return True
+
+        # Rule 4 — FloorDiv(a, b): if a is multiple of b*n
+        if isinstance(numerator, FloorDiv):
+            a, b = numerator.args
+            if isinstance(b, (int, sympy.Integer)):
+                if self._is_multiple_of(a, int(b) * denominator):
+                    return True
+
+        # Rule 5 — Mod(a, b): Mod(a,b) = a - b*floor(a/b), so if both a and b
+        # are multiples of n, then Mod(a,b) is too.
+        if isinstance(numerator, (Mod, sympy.Mod)):
+            a, b = numerator.args
+            if self._is_multiple_of(a, denominator) and self._is_multiple_of(
+                b, denominator
+            ):
+                return True
+
+        # Rule 6 — axiom fallback: ask ShapeEnv
+        expr = sympy.Eq(Mod(numerator, denominator), 0)
+        return self.statically_known_true(expr)
+
+    def statically_known_multiple_of(
+        self, numerator: Expr, denominator: Expr | int
+    ) -> bool:
+        """
+        Return a bool indicating if it is sound to optimize for the numerator being a multiple of the denominator.
+        """
+        # The reason we skip compute here is to avoid the cost of trying to eval this symbolically.
+        # see https://github.com/sympy/sympy/issues/28200
+
+        if len(free_symbols(numerator)) > 20:
+            return False
+
+        if isinstance(denominator, (int, sympy.Integer)):
+            return self._is_multiple_of(numerator, int(denominator))
+
+        # For symbolic denominators, fall back to direct sympy check
+        expr = sympy.Eq(Mod(numerator, denominator), 0)
+        return self.statically_known_true(expr)  # type: ignore[arg-type]
+
+    def statically_known_power_of_2(self, expr: Expr) -> bool:
+        """
+        Returns a bool indicating if x is known to be a power of 2.
+        """
+        return isinstance(expr, sympy.Integer) and is_power_of_2(int(expr))
+
+    # The expect/check functions require you to ALREADY KNOW that a particular
+    # condition holds. They are similar to expect_true in symbolic_shapes.py and
+    # torch.check but operates on sympy expressions instead of symnodes.
+    def expect_true(self, expr: Expr) -> bool:
+        """
+        Use it when you already know that expr is true or should be true and want to
+        ensure that guards/runtime assertions are in place to ensure this in compiled
+        function. Unlike check, this WON'T raise an error if expr isn't actually true.
+        check Note [expect_true].
+        """
+        if not self.statically_known_true(expr):
+            return self.shape_env.guard_or_defer_runtime_assert(
+                expr, "sizevars.expect_true"
+            )
+        return True
+
+    def check(self, expr: Expr) -> None:
+        """
+        Use it when you already know that expr is true or should be true and want to
+        ensure that guards/runtime assertions are in place to ensure this in compiled
+        function. Unlike expect_true, this WILL raise an error if expr isn't actually true.
+        check Note [expect_true].
+        """
+        expr = sympy_subs(expr, self.inv_precomputed_replacements)
+        assert self.expect_true(expr)
+
+    def check_equals(self, left: Expr, right: Expr) -> None:
+        """
+        check(sympy.Eq(left, right)).
+
+        """
+        self.check(sympy.Eq(left, right))
+        return left
+
+    def check_equals_and_simplify(self, left: Expr, right: Expr) -> Expr:
+        """
+        check(sympy.Eq(left, right)) and returns left after applying
+        inv_precomputed_replacements.
+        """
+        self.check(sympy.Eq(left, right))
+        return sympy_subs(left, self.inv_precomputed_replacements)
+
+    def check_leq(self, left: Expr, right: Expr) -> None:
+        self.check(sympy.Le(left, right))
+
+    def check_lt(self, left: Expr, right: Expr) -> None:
+        self.check(sympy.Lt(left, right))
+
+    # Similar to the functions guard_or_false/guard_or_true in symbolic_shapes.py
+    # but operates on sympy expressions instead of symnodes. see Note [guard_or_].
+    def guard_or_false(self, left):
+        import torch.fx.experimental._config as exp_config
+
+        if exp_config.backed_size_oblivious:
+            static_val = self.shape_env._maybe_evaluate_static(left)
+            if static_val is not None:
+                return static_val
+            return False
+        return self.evaluate_expr(left, fallback_value=False)
+
+    def guard_or_true(self, left):
+        import torch.fx.experimental._config as exp_config
+
+        if exp_config.backed_size_oblivious:
+            static_val = self.shape_env._maybe_evaluate_static(left)
+            if static_val is not None:
+                return static_val
+            return True
+        return self.evaluate_expr(left, fallback_value=True)
+
+    # The evaluate functions evaluate some symbolic sympy expression
+    # (NB: not necessarily an Expr) and return what the concrete result
+    # is, guarding on the expression being that result
+
+    # NB: write evaluate_expr(sympy.Lt(a, b)) rather than evaluate_expr(a < b)
+    # as this will ensure that you actually have a sympy'ified expression,
+    # and will prevent you from incorrectly writing evaluate_expr(a == b)
+    # which does the wrong thing if a or b is a sympy expression
+    def evaluate_expr(
+        self,
+        left: Expr | sympy.logic.boolalg.Boolean,
+        size_oblivious: bool = False,
+        fallback_value: bool | None = None,
+    ) -> bool:
+        assert isinstance(left, (Expr, sympy.logic.boolalg.Boolean)), type(left)
+        return self.shape_env.evaluate_expr(
+            sympy.sympify(left),
+            size_oblivious=size_oblivious,
+            fallback_value=fallback_value,
+        )
+
+    def is_size_one_or_false(self, size: Expr) -> bool:
+        """Return True if size equals 1.
+
+        Unbacked symbolic sizes return False without introducing a guard.
+        """
+        return self.guard_or_false(sympy.Eq(size, 1))
+
+    def evaluate_min(self, left: Expr, right: Expr) -> Expr:
+        """Return the smaller of left and right, and guard on that choice."""
+        if isinstance(left, Expr):
+            left = sympy_subs(left, self.inv_precomputed_replacements)  # type: ignore[arg-type]
+        if isinstance(right, Expr):
+            right = sympy_subs(right, self.inv_precomputed_replacements)  # type: ignore[arg-type]
+        if self.guard_or_false(sympy.Le(left, right)):
+            return left
+        if self.guard_or_false(sympy.Le(right, left)):
+            return right
+
+        # GCD fallback: if gcd(a, b) == a then a divides b, implying a <= b.
+        #
+        # TODO: This is NOT always sound for unbacked symints.  It can
+        # produce wrong results when:
+        #   - inputs can be negative: gcd(u0, 10*u0) = u0, returns u0,
+        #     but if u0 < 0 then u0 > 10*u0 (e.g. u0=-1: min(-1,-10) = -10)
+        #   - a factor can be zero: gcd(u0, u0*u1) = u0, returns u0,
+        #     but if u1=0 then u0*u1=0 < u0 (e.g. u0=5,u1=0: min(5,0) = 0)
+        # TODO shall we add a runtime assertion at least.
+        gcd = sympy.gcd(left, right)
+        if left == gcd:
+            return left
+        if right == gcd:
+            return right
+
+        # Min/Max fallback: we can prove Min(a, b) <= c when any arg <= c, but
+        # sympy doesn't simplify this yet. So, evaluate it here. Same for Max.
+        for lhs, rhs in [(left, right), (right, left)]:
+
+            def le_rhs(a: Expr) -> bool:
+                return self.guard_or_false(sympy.Le(a, rhs))
+
+            # Min(Min(a, b), c) ==> Min(a, b) if (a <= c) or (b <= c).
+            if isinstance(lhs, sympy.Min) and any(le_rhs(a) for a in lhs.args):
+                return lhs
+            # Min(Max(a, b), c) ==> Max(a, b) if (a <= c) and (b <= c).
+            if isinstance(lhs, sympy.Max) and all(le_rhs(a) for a in lhs.args):
+                return lhs
+
+        raise TypeError(
+            f"evaluate_min({left}, {right}) with unbacked symints"
+        ) from None
+
+    def evaluate_max(self, left: Expr, right: Expr) -> Expr:
+        """Return the larger of left and right, and guard on that choice."""
+        # Always choose the opposite of eval min for consistency
+        # This means min(a, b) and max(a, b) produce the same guards
+        min_val = self.evaluate_min(left, right)
+        return right if min_val is left else left
+
+    def guard_int(self, expr: Expr | int) -> int:
+        """
+        Similar to guard_int in symbolic_shapes.py, except this function works with SymPy
+        expressions instead of SymNodes. It extracts the value represented by expr from shapeEnv
+        and specialize the compiled graph on it. Raises an error if the result cannot be
+        determined due to unhinted or unbacked symbols.
+        """
+        if isinstance(expr, int):
+            return expr
+        val = self.guarding_hint_or_throw(expr)
+        self.check_equals(expr, sympy.Integer(val))
+        return int(val)
+
+    def guard_int_seq(self, left: Sequence[Expr | int]) -> list[int]:
+        """
+        Apply guard_int on a sequence of inputs.
+        """
+        return [self.guard_int(x) for x in left]
+
+    def remove_precomputed_replacements(self, expr: Expr) -> Expr:
+        if any(symbol_is_type(s, SymT.PRECOMPUTED_SIZE) for s in expr.free_symbols):  # type: ignore[attr-defined]
+            return sympy_subs(expr, self.inv_precomputed_replacements)  # type: ignore[arg-type]
+        return expr
+
+    def replace_backed_symbols_with_hints(
+        self,
+        expr: Expr | int,
+    ) -> Expr | int:
+        """
+        Replace all backed symbols in an expression with their concrete hint values.
+
+        This function substitutes backed symbolic variables with their known concrete
+        values from the shape environment, while leaving unbacked symbols (data-dependent
+        values) untouched. Not recommended to be used except as last resort!
+        was added for existing usages compatibility.
+        """
+        if isinstance(expr, int):
+            return expr
+        # Substitute all hints into expr, but leave unbacked symints alone
+        expr = self.simplify(expr)
+        if not isinstance(expr, Expr):
+            assert isinstance(expr, int)
+            return expr
+
+        expr = self.remove_precomputed_replacements(expr)
+        expr = sympy_subs(expr, self.backed_var_to_val)
+        if isinstance(expr, Expr):
+            expr = expr.expand(identity=True)
+
+        free_symbols = expr.free_symbols
+        if not free_symbols:
+            try:
+                return int(expr)
+            except TypeError:
+                return expr  # inf/nan/I
+
+        return expr
+
+    def to_symint_or_int(self, expr: Expr | int) -> SymInt | int:
+        """Convert a sympy expression to SymInt, or return int as is."""
+        if isinstance(expr, int):
+            return expr
+
+        from torch.fx.experimental.sym_node import SymNode
+
+        expr = self.simplify(expr)
+        if isinstance(expr, (int, sympy.Integer)):
+            return int(expr)
+
+        # Remove precomputed replacements to get canonical form
+        # (consistent with symbolic_hint behavior)
+        expr = self.remove_precomputed_replacements(expr)
+
+        try:
+            hint = self.guarding_hint_or_throw(expr)
+        except Exception:
+            hint = None
+        node = SymNode(expr, self.shape_env, int, hint)
+        return SymInt(node)
+
+    def to_symints_or_ints(self, exprs: Sequence[Expr | int]) -> list[SymInt | int]:
+        """Convert a sequence of sympy expressions to SymInts, or return ints as is."""
+        return [self.to_symint_or_int(e) for e in exprs]
+
+    def guarding_hint_or_throw(self, expr: Expr | int) -> int:
+        """
+        Return a concrete integer hint for an expression that is safe to use for guarding.
+
+        This method evaluates the expression using only backed-symbols hints. Unlike
+        optimization_hint(), this method does NOT use heuristics or fallback values
+        for unbacked symbols.
+
+        Use this method when you need a hint value that will be used for guarding decision.
+
+        Args:
+            expr: A sympy expression or integer to evaluate.
+
+        Returns:
+            The concrete integer value of the expression based on backed symbol hints.
+
+        Raises:
+            GuardOnDataDependentSymNode: If the expression contains unbacked symbols
+            (data-dependent values) that cannot be resolved to concrete values.
+
+        See Also:
+            optimization_hint: For cases where fallback/heuristic values are acceptable
+                for unbacked symbols.
+        """
+        if isinstance(expr, SymNode):
+            raise TypeError(
+                f"guarding_hint_or_throw expects a sympy Expr or int, not {type(expr)}. "
+                "Use expr.expr to extract the sympy expression from a SymNode."
+            )
+        return _guarding_hint_or_throw_base(
+            self.shape_env, expr, self.inv_precomputed_replacements
+        )
+
+    def optimization_hint(self, expr: Expr | int, fallback: int | None = None) -> int:
+        """
+        Return a concrete integer hint for an expression.
+
+        This function should be used for non-guarding based optimizations. If you
+        want a hint that you can guard on, use the guarding_hint API instead.
+
+        This function will hint unbacked symbols using user provided optimization
+        hints. If not provided, fallback will be used along with some heuristics
+        that try to maximize consistency with the shape environment.
+
+        Special cases:
+        - Complex numbers (containing sympy.I): raises an error since tensor
+          dimensions cannot be complex.
+        - Infinity (int_oo, sympy.oo): returns sys.maxsize.
+        - NaN (sympy.nan): returns the fallback value.
+        """
+        return _optimization_hint_base(
+            self.shape_env, expr, self.inv_precomputed_replacements, fallback
+        )
+
+    def optimization_hints(
+        self,
+        exprs: Iterable[Expr | int],
+        fallback: int | None = None,
+    ) -> tuple[int, ...]:
+        """
+        Like optimization_hint but for a sequence of expressions.
+        Returns a tuple of concrete integer hints.
+        """
+        # Read config at call time to respect runtime patches (e.g., in tests)
+        if fallback is None:
+            fallback = config.unbacked_symint_fallback
+        return tuple(self.optimization_hint(x, fallback=fallback) for x in exprs)
+
+    def all_unbacked_explicitly_hinted(self, exprs: IterateExprs) -> bool:
+        """
+        Return True if every unbacked symbol in *exprs* has an explicit
+        user-provided hint in var_to_hint_override.  If there are no
+        unbacked symbols at all, returns True (vacuously).
+        *exprs* can be a single expression or any iterable accepted by
+        free_unbacked_symbols (list, tuple, etc.).
+        """
+        unbacked = free_unbacked_symbols(exprs)
+        return unbacked.issubset(self.var_to_hint_override.keys())
+
+    def optimization_hint_with_override(
+        self,
+        expr: Expr | int,
+        hint_override: int | None,
+    ) -> int:
+        r"""Return a concrete integer hint for an expression, with optional override.
+        This is used in dynamic dispatch scenarios where callers may want to
+        provide a specific hint value rather than computing one from the expression.
+        The resolution order is:
+        1. If ``expr`` simplifies to a static integer, return that value
+           (``hint_override`` is ignored for static shapes).
+        2. If ``expr`` is dynamic and ``hint_override`` is not ``None``,
+           return ``hint_override``.
+        3. Otherwise, compute a hint via :meth:`optimization_hint`
+        Args:
+            expr (Expr or int): The expression to get a hint for.
+            hint_override (int, optional): If provided and ``expr`` is dynamic,
+                this value is returned instead of computing a hint.
+        Returns:
+            int: A concrete integer hint for the expression.
+        """
+        simplified = _maybe_realize_expr(self.simplify(expr), None)
+
+        if simplified is not None:
+            return simplified
+
+        # Dynamic shape: use hint_override if set,
+        # else return optimization_hint
+        if hint_override is not None:
+            return hint_override
+
+        return self.optimization_hint(expr)
+
+    def optimization_hints_with_override(
+        self,
+        exprs: Iterable[Expr | int],
+        hint_override: int | None,
+    ) -> tuple[int, ...]:
+        """
+        Like optimization_hint_with_override but for a sequence of expressions.
+        Returns a tuple of concrete integer hints.
+        """
+        return tuple(
+            self.optimization_hint_with_override(e, hint_override) for e in exprs
+        )
+
+    def guarding_hints_or_throw(
+        self,
+        exprs: Iterable[Expr | int],
+    ) -> tuple[int, ...]:
+        return tuple(self.guarding_hint_or_throw(x) for x in exprs)
+
+    def _lru_cache(self, fn, maxsize=None):
+        """
+        Wrapper around functools.lru_cache that clears when replacements
+        has been invalidated.
+        """
+        fn_cache = functools.lru_cache(maxsize)(fn)
+        prior_len = len(self.replacements)
+
+        @functools.wraps(fn)
+        def wrapper(*args, **kwargs):
+            nonlocal prior_len
+            if prior_len != len(self.replacements):
+                prior_len = len(self.replacements)
+                fn_cache.cache_clear()
+            return fn_cache(*args, **kwargs)
+
+        return wrapper
+
+    def make_stride_vars_cache(self):
+        cache = self._lru_cache(self._stride_vars)
+
+        def stride_vars(
+            index: Expr,
+            vars: Sequence[sympy.Symbol],
+            support_vars: Sequence[sympy.Symbol] | None = None,
+        ) -> list[Expr]:
+            if not support_vars:
+                support_vars = vars
+            return cache(index, tuple(vars), tuple(support_vars))
+
+        return stride_vars
+
+    def _stride_vars(
+        self,
+        index: Expr,
+        vars: Sequence[sympy.Symbol],
+        support_vars: Sequence[sympy.Symbol],
+    ) -> list[Expr]:
+        """Convert an indexing expression back into strides
+
+        NOTE: This is only valid if the index is a standard strided offset
+        calculation. e.g. 10 * ModularIndexing(i0 + 1, 1, 2) would give a
+        stride of -10 because the index wraps around after the first element
+
+        """
+        strides = []
+        index = self.simplify(index)
+        # remove any offset
+        index = index - sympy_subs(
+            index, {v: sympy.S.Zero for v in support_vars if v != 0}
+        )
+        for i in range(len(vars)):
+            # drop all the other dims
+            index_dim = sympy_subs(
+                index,
+                {
+                    support_vars[j]: sympy.S.Zero
+                    for j in range(len(support_vars))
+                    if vars[i] != support_vars[j] and support_vars[j] != 0
+                },
+            )
+            v = vars[i]
+            if v == 0:
+                strides.append(sympy.S.Zero)
+            else:
+                # TODO(jansel): should we use sympy.diff here?
+                strides.append(
+                    sympy_subs(index_dim, {v: sympy.S.One})
+                    - sympy_subs(index_dim, {v: sympy.S.Zero})
+                )
+        return strides
+
+    def offset_var(self, index: Expr, vars: Sequence[sympy.Symbol]) -> Expr:
+        """Extract offset part of an indexing expression"""
+        index = self.simplify(index)
+        return sympy_subs(index, {v: sympy.S.Zero for v in vars if v != 0})
+
+    # Return stride optimizaitons hints,
+    # only used for optimizations.
+    def stride_hints(
+        self,
+        index: Expr,
+        vars: Sequence[sympy.Symbol],
+        support_vars: Sequence[sympy.Symbol] | None = None,
+    ) -> list[int]:
+        for v in index.free_symbols:
+            if symbol_is_type(v, SymT.INDIRECT):  # type: ignore[attr-defined]
+                index = sympy_subs(index, {v: 0})  # type: ignore[dict-item]
+        result = []
+        for s in self.stride_vars(index, vars, support_vars):
+            result.append(self.optimization_hint(s, fallback=0))
+        return result
+
+    def stride_order(self, index: Expr, vars: list[sympy.Symbol]) -> list[int]:
+        strides = tuple(map(abs, self.stride_hints(index, vars)))
+        order = list(range(len(strides)))
+        order.sort(key=lambda x: (strides[x] == 0, strides[x]))
+        return order
+
+    def lookup_precomputed_size(self, expr: Expr) -> Expr:
+        if (
+            isinstance(expr, (int, sympy.Symbol, sympy.Number))
+            or expr.is_number
+            or expr.is_symbol
+        ):
+            return expr
+        expr = self.remove_precomputed_replacements(expr)
+        if expr not in self.precomputed_replacements:
+            sym = sympy_index_symbol_with_prefix(
+                SymT.PRECOMPUTED_SIZE, len(self.precomputed_replacements)
+            )
+            self.precomputed_replacements[expr] = sym
+            self.inv_precomputed_replacements[sym] = expr
+        return self.precomputed_replacements[expr]
+
+    def free_symbols(self) -> OrderedSet[sympy.Symbol]:
+        return OrderedSet(self.backed_var_to_val.keys()) - OrderedSet(
+            self.replacements.keys()
+        )
+
+    def combine_modular_indexing_pairs(self, index: sympy.Expr) -> sympy.Expr:
+        """
+        A pair of special ModularIndexing can be combined.
+
+        E.g. ModularIndexing(ModularIndexing(x, 1, a), 1, b)
+        We can simplify this to ModuleIndexing(x, 1, b), if
+        1. x is non negative integer
+        2. a and b are positive integers
+        3. a is a multiple of b.
+        """
+
+        def _check_args(x, div, mod, is_first):
+            if not isinstance(div, sympy.Integer) or not isinstance(mod, sympy.Integer):
+                return False
+            if div != 1:
+                return False
+            if mod <= 0:
+                return False
+
+            if is_first:
+                # first ModularIndexing should contains a nested ModularIndex
+                if not isinstance(x, ModularIndexing):
+                    return False
+            else:
+                # second ModularIndexing should contains a non-negative
+                # symbol
+                if not isinstance(x, sympy.Symbol) or not self.statically_known_geq(
+                    x, 0
+                ):
+                    return False
+            return True
+
+        if isinstance(index, ModularIndexing):
+            x, div, mod = index.args
+
+            if not _check_args(x, div, mod, True):
+                return index
+
+            x2, div2, mod2 = x.args
+
+            if not _check_args(x2, div2, mod2, False):
+                return index
+
+            if Mod(mod2, mod) != 0:
+                return index
+
+            return ModularIndexing(x2, 1, mod)
+
+        return index
+
+    def expand_floor_div(
+        self, index: sympy.Expr
+    ) -> bool | tuple[sympy.Expr, sympy.Expr]:
+        """
+        Expand the FloorDiv to the entire expression so that the expression may
+        be simplified.
+
+        E.g., for a 2D contiguous tensor with shape [a, 2 * b], and index variables
+        x1, x2, index expression 'x1 * 2b + x2' can be easily combined.
+        But index expression 'x1 * b + x2 // 2' can not.
+        By expanding the FloorDiv to the entire expression, we get
+        '(x1 * 2b + x2) // 2'. This transformation allows us to merge loops
+        for the numerator!
+
+        Return false if this optimization can be applied;
+        Return the new expression and the denominator otherwise.
+        The original expression will be equivalent to 'new_expression // denominator'
+        """
+        if not isinstance(index, sympy.Add):
+            return False
+        terms = index.args
+
+        if len(terms) < 2:
+            return False
+        floor_div_index = -1
+        varlist = []
+        factorlist = []
+        for idx, term in enumerate(terms):
+            if isinstance(term, sympy.Mul):
+                # For dynamic shape, term like '2*s1*x1' has 3 child nodes.
+                # - A integer for 2
+                # - A symbol for s1
+                # - A symbol for x1
+                # Skip for now.
+                if len(term.args) != 2:
+                    return False
+                factor, var = term.args
+                varlist.append(var)
+                factorlist.append(factor)
+                if not isinstance(factor, sympy.Integer) or not isinstance(
+                    var, sympy.Symbol
+                ):
+                    return False
+                # It's easier to reason about the correceness of the transformation
+                # for non-negative integers.
+                if not self.statically_known_geq(var, 0):
+                    return False
+            elif isinstance(term, FloorDiv):
+                var, factor = term.args
+                if not isinstance(factor, sympy.Integer) or not isinstance(
+                    var, sympy.Symbol
+                ):
+                    return False
+                if not self.statically_known_geq(var, 0):
+                    return False
+                if floor_div_index >= 0:
+                    # can not handle multi FloorDiv yet
+                    return False
+
+                floor_div_index = idx
+                varlist.append(var)
+                # this factor is denominator
+                factorlist.append(factor)
+            else:
+                return False
+
+        if floor_div_index < 0:
+            return False
+
+        # Construct the new expression and remember the denominator
+        denominator = factorlist[floor_div_index]
+        new_index = sympy.S.Zero
+
+        for var, factor, idx in zip(varlist, factorlist, itertools.count()):
+            if idx == floor_div_index:
+                new_index += var
+            else:
+                new_index += (factor * denominator) * var
+
+        return new_index, denominator
+
+
+def join_dimensions(expr: Expr) -> Expr:
+    if not isinstance(expr, sympy.Add) or not expr.has(ModularIndexing):
+        return expr  # fast exit path
+    return _join_dimensions_cached(expr)
+
+
+@functools.lru_cache(256)
+def _join_dimensions_cached(expr: Expr) -> Expr:
+    """
+    ModularIndexing(i0, 1, 32) + 32 * ModularIndexing(i0, 32, 4)
+    becomes
+    ModularIndexing(i0, 1, 128)
+    ModularIndexing(i0, 1, 32) + 32 * FloorDiv(i0, 32)
+    becomes i0
+
+
+    This type of pattern can come from view operations
+    """
+    assert isinstance(expr, sympy.Add)
+
+    scale = sympy.Wild("scale", exclude=[0], integer=True)
+    base = sympy.Wild("base", integer=True)
+    divisor = sympy.Wild("divisor", integer=True)
+    mod1 = sympy.Wild("modulus", integer=True)
+    mod2 = sympy.Wild("modulus2", integer=True)
+    for term1 in expr.args:
+        m1 = term1.match(scale * ModularIndexing(base, divisor, mod1))
+        if m1:
+            for term2 in expr.args:
+                m2 = term2.match(
+                    m1[scale]
+                    * m1[mod1]
+                    * ModularIndexing(m1[base], m1[divisor] * m1[mod1], mod2)
+                )
+                if m2 and term1 != term2:
+                    expr = join_dimensions(
+                        expr
+                        - term1
+                        - term2
+                        + m1[scale]
+                        * ModularIndexing(m1[base], m1[divisor], m1[mod1] * m2[mod2])
+                    )
+                    return expr
+    for term1 in expr.args:
+        m1 = term1.match(scale * ModularIndexing(base, divisor, mod1))
+        if m1:
+            for term2 in expr.args:
+                m2 = term2.match(
+                    m1[scale] * m1[mod1] * FloorDiv(m1[base], m1[divisor] * m1[mod1])
+                )
+                if m2 is not None:  # in case of success we get an empty dict here
+                    expr = join_dimensions(
+                        expr
+                        - term1
+                        - term2
+                        + m1[scale] * FloorDiv(m1[base], m1[divisor])
+                    )
+                    return expr
+    return expr
+
+
+class SimplifyIndexing(V.WrapperHandler):  # type: ignore[name-defined]
+    """
+    A wrapper around .virtualize.ops that uses var range information to
+    simplify ModularIndexing/FloorDiv.
+    """
+
+    def __init__(self, inner, var_ranges: VarRanges) -> None:
+        super().__init__(inner)
+        self.name = "SimplifyIndexing"
+        self._simplify: Callable[[Expr], Expr] = (
+            lambda index: V.graph.sizevars.simplify_with_ranges(index, var_ranges)
+        )
+
+    def load(self, name: str, index: sympy.Expr):
+        return self._inner.load(name, self._simplify(index))
+
+    def store(self, name, index, value, mode=None):
+        return self._inner.store(name, self._simplify(index), value, mode=mode)
+
+    def store_reduction(self, name, index, value):
+        return self._inner.store_reduction(name, self._simplify(index), value)
+
+    def index_expr(self, index, dtype):
+        return self._inner.index_expr(self._simplify(index), dtype)
+
+    def check_bounds(self, index, size, lower, upper):
+        return self._inner.check_bounds(self._simplify(index), size, lower, upper)

outputs/audit_venv/lib/python3.11/site-packages/torch/_inductor/standalone_compile.py

@@ -0,0 +1,491 @@
+from __future__ import annotations
+
+import contextlib
+import copy
+import logging
+import os
+import pickle
+import shutil
+from abc import ABC, abstractmethod
+from contextlib import AbstractContextManager, nullcontext
+from typing import Any, Literal, TYPE_CHECKING
+
+import torch.fx
+from torch._dynamo.aot_compile_types import BundledAOTAutogradSerializableCallable
+from torch._dynamo.utils import dynamo_timed
+from torch._inductor.cpp_builder import normalize_path_separator
+from torch._inductor.cudagraph_utils import BoxedDeviceIndex
+from torch._inductor.runtime.cache_dir_utils import temporary_cache_dir
+from torch._inductor.utils import BoxedBool, InputType
+from torch._subclasses import FakeTensorMode
+from torch.fx.experimental.symbolic_shapes import ShapeEnv
+
+from . import config
+
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Sequence
+
+    from torch.compiler._cache import CacheInfo
+    from torch.fx import GraphModule
+
+
+log = logging.getLogger(__name__)
+
+
+class CompiledArtifact(ABC):
+    """
+    CompiledArtifact class represents the inductor cache artifacts that
+    can be invoked in order to avoid repeated compilation.
+
+    CompiledArtifact can be obtained by calling standalone_compile(gm, example_inputs)
+    to create a fresh CompiledArtifact from a GraphModule and example inputs.
+
+    Later this CompiledArtifact can be saved to disk, either as a binary or unpacked
+    into the provided folder via the CompiledArtifact.save function.
+
+    CompiledArtifact.load provides a way to create a CompiledArtifact from the
+    binary or unpacked data.
+
+    Finally, the CompiledArtifact can be invoked via the __call__ method
+    to execute the cached artifact.
+    """
+
+    def __init__(
+        self,
+        compiled_fn: Callable[..., Any],
+        artifacts: tuple[bytes, CacheInfo] | None,
+    ):
+        self._compiled_fn = compiled_fn
+        self._artifacts = artifacts
+
+    @abstractmethod
+    def __call__(self, *args: Any) -> Any: ...
+
+    @abstractmethod
+    def save(
+        self, *, path: str, format: Literal["binary", "unpacked"] = "binary"
+    ) -> None: ...
+
+    @staticmethod
+    def load(
+        *, path: str, format: Literal["binary", "unpacked"] = "binary"
+    ) -> CompiledArtifact:
+        if format == "unpacked":
+            # If format is unpacked, it must be a CacheCompiledArtifact
+            return CacheCompiledArtifact.load(path=path, format=format)
+
+        assert format == "binary"
+        with open(path, "rb") as file:
+            from torch.utils._appending_byte_serializer import BytesReader
+
+            from .codecache import torch_key
+
+            result_bytes = file.read()
+            reader = BytesReader(result_bytes)
+            header = reader.read_bytes()
+            if header == AOTCompiledArtifact.AOT_HEADER:
+                assert reader.read_bytes() == torch_key()
+                artifact = reader.read_bytes()
+                assert reader.is_finished()
+                return AOTCompiledArtifact.deserialize(artifact)
+            # Otherwise, it's in the CacheCompiledArtifact format
+            elif header == CacheCompiledArtifact.CACHE_HEADER:
+                assert reader.read_bytes() == torch_key()
+                key = reader.read_str()
+                artifact_bytes = reader.read_bytes()
+                assert reader.is_finished()
+                torch.compiler.load_cache_artifacts(artifact_bytes)
+                return CacheCompiledArtifact._load_impl(nullcontext(), key)
+            else:
+                raise RuntimeError(
+                    "Invalid header, expected CacheCompiledArtifact or AOTCompiledArtifact, got: "
+                    + header.decode("utf-8")
+                )
+
+
+class CacheCompiledArtifact(CompiledArtifact):
+    """
+    CompiledArtifact that depends on torch.compiler.save_cache_artifacts
+    """
+
+    CACHE_HEADER = bytes("CacheCompiledArtifact", "utf-8")
+
+    def __init__(
+        self,
+        compiled_fn: Callable[..., Any],
+        artifacts: tuple[bytes, CacheInfo] | None,
+    ):
+        self._compiled_fn = compiled_fn
+        self._artifacts = artifacts
+
+    def __call__(self, *args: Any) -> Any:
+        return self._compiled_fn(*args)
+
+    def is_saveable(self) -> bool:
+        if self._artifacts is None:
+            return False
+        _, cache_info = self._artifacts
+        # 0 means nothing was saved
+        # >1 means multiple artifacts were saved, which is concerning
+        # (we only expect one)
+        return len(cache_info.aot_autograd_artifacts) == 1
+
+    def save(
+        self, *, path: str, format: Literal["binary", "unpacked"] = "binary"
+    ) -> None:
+        with dynamo_timed("CompiledArtifact.save"):
+            if self._artifacts is None:
+                raise RuntimeError(
+                    "CompiledArtifact.save failed to save since there's no artifact to save"
+                )
+            artifact_bytes, cache_info = self._artifacts
+            if len(cache_info.aot_autograd_artifacts) == 0:
+                raise RuntimeError(
+                    f"CompiledArtifact.save failed to save due to no aot_autograd artifacts. "
+                    f"This likely means there was something that was not serializable in the "
+                    f"graph passed to standalone_compile. This can generally be fixed by "
+                    f"ensuring that your model only uses constructs that are serializable. "
+                    f"{cache_info}"
+                )
+            if len(cache_info.aot_autograd_artifacts) > 1:
+                raise AssertionError(
+                    f"CompiledArtifact.save failed to save because there was more than one "
+                    f"artifact but we only expected one. {cache_info}"
+                )
+            key = cache_info.aot_autograd_artifacts[0]
+
+            if format == "binary":
+                # can't assert that it is a file since it might not exist yet
+                assert not os.path.isdir(path)
+
+                from torch.utils._appending_byte_serializer import BytesWriter
+
+                from .codecache import torch_key
+
+                writer = BytesWriter()
+                writer.write_bytes(CacheCompiledArtifact.CACHE_HEADER)
+                writer.write_bytes(torch_key())
+                writer.write_str(key)
+                writer.write_bytes(artifact_bytes)
+
+                from torch._inductor.codecache import write_atomic
+
+                write_atomic(path, writer.to_bytes())
+            else:
+                assert format == "unpacked"
+                if os.path.exists(path):
+                    assert os.path.isdir(path)
+                    shutil.rmtree(path, ignore_errors=True)
+
+                from .codecache import FxGraphCache
+
+                with temporary_cache_dir(path):
+                    # This function unpacks the cache artifacts to disk
+                    loaded_cache_info = torch.compiler.load_cache_artifacts(
+                        artifact_bytes
+                    )
+                    assert loaded_cache_info is not None
+                    # Now write all the output_code artifacts to disk so that
+                    # they can be inspected and modified
+                    for key in loaded_cache_info.inductor_artifacts:
+                        subdir = FxGraphCache._get_tmp_dir_for_key(key)
+                        assert os.path.exists(subdir)
+                        for path in sorted(os.listdir(subdir)):
+                            with open(os.path.join(subdir, path), "rb") as f:
+                                graph = pickle.load(f)
+                            output_file = graph.write_to_disk()
+                            log.info("Output code written to: %s", output_file)
+
+    @staticmethod
+    def _load_impl(
+        cache_dir_ctx: AbstractContextManager[Any], key: str
+    ) -> CompiledArtifact:
+        with (
+            cache_dir_ctx,
+            config.patch(unsafe_skip_cache_dynamic_shape_guards=True),
+        ):
+            with torch._functorch.config.patch(strict_autograd_cache=True):
+                from torch._functorch._aot_autograd.autograd_cache import (
+                    AOTAutogradCache,
+                )
+
+                result = AOTAutogradCache._lookup(
+                    key,
+                    local=True,
+                    remote=False,
+                    args=[],
+                    cache_info={},
+                    aot_config=None,
+                )
+
+            assert result is not None
+            (entry, _) = result
+
+            from .compile_fx import _CompileFxKwargs
+
+            fx_config = _CompileFxKwargs(
+                cudagraphs=BoxedBool(False),
+                boxed_forward_device_index=BoxedDeviceIndex(0),
+            )
+
+            context = torch._guards.TracingContext(FakeTensorMode(shape_env=ShapeEnv()))
+            with torch._guards.tracing(context):
+                compiled_fn = entry.wrap_post_compile(
+                    [], entry.sanitized_aot_config, fx_config
+                )
+        return CacheCompiledArtifact(lambda *args: compiled_fn(list(args)), None)
+
+    @staticmethod
+    def _prepare_load(
+        *, path: str, format: Literal["binary", "unpacked"] = "binary"
+    ) -> tuple[str, AbstractContextManager[Any]]:
+        """
+        Do format specific prep and loads, return a context manager and key
+        """
+        path = normalize_path_separator(path)
+        with dynamo_timed("CompiledArtifact.load"):
+            if format == "binary":
+                # can't assert that it is a file since it might not exist yet
+                assert not os.path.isdir(path)
+                with open(path, "rb") as file:
+                    artifacts = file.read()
+                from torch.utils._appending_byte_serializer import BytesReader
+
+                from .codecache import torch_key
+
+                reader = BytesReader(artifacts)
+                assert reader.read_bytes() == torch_key()
+                key = reader.read_str()
+                artifact_bytes = reader.read_bytes()
+                assert reader.is_finished()
+
+                torch.compiler.load_cache_artifacts(artifact_bytes)
+                return key, nullcontext()
+            else:
+                assert format == "unpacked"
+                assert os.path.isdir(path)
+                autograd_cache_dir = os.path.join(path, "aotautograd")
+                assert os.path.isdir(autograd_cache_dir)
+                files = list(os.listdir(autograd_cache_dir))
+                assert len(files) == 1
+                key = files[0]
+                cache_dir_ctx = temporary_cache_dir(path)
+                return key, cache_dir_ctx
+
+    @staticmethod
+    def load(
+        *, path: str, format: Literal["binary", "unpacked"] = "binary"
+    ) -> CompiledArtifact:
+        key, cache_dir_ctx = CacheCompiledArtifact._prepare_load(
+            path=path, format=format
+        )
+        return CacheCompiledArtifact._load_impl(cache_dir_ctx, key)
+
+
+class AOTCompiledArtifact(CompiledArtifact):
+    """
+    Similar to CompiledArtifact, but the object is a single, bundled precompiled function.
+    This object is always a serializable callable function.
+
+    This object is essentially a wrapper for BundledAOTAutogradSerializableCallable, which
+    is used by torch._dynamo.aot_compile for AOT Precompilation.
+    """
+
+    AOT_HEADER = bytes("AOTCompiledArtifact", "utf-8")
+
+    def __init__(
+        self,
+        compiled_fn: Callable[..., Any],
+    ):
+        self.inner_fn = BundledAOTAutogradSerializableCallable(compiled_fn)
+        self._artifacts = (
+            None  # We don't need artifacts, the inner object handles everything
+        )
+
+    @staticmethod
+    def from_bundled_callable(
+        bundled_fn: BundledAOTAutogradSerializableCallable,
+    ) -> AOTCompiledArtifact:
+        return AOTCompiledArtifact(bundled_fn.compiled_fn)
+
+    def __call__(self, *args: Any) -> Any:
+        return self.inner_fn(*args)
+
+    def save(
+        self, *, path: str, format: Literal["binary", "unpacked"] = "binary"
+    ) -> None:
+        if format == "unpacked":
+            raise RuntimeError(
+                "AOTCompiledArtifact does not support unpacked format yet"
+            )
+        result_bytes = self.serialize()
+        from torch.utils._appending_byte_serializer import BytesWriter
+
+        from .codecache import torch_key
+
+        writer = BytesWriter()
+        writer.write_bytes(AOTCompiledArtifact.AOT_HEADER)
+        writer.write_bytes(torch_key())
+        writer.write_bytes(result_bytes)
+
+        from torch._inductor.codecache import write_atomic
+
+        # Save a sentinel file to indicate that this is AOT
+        write_atomic(path, writer.to_bytes())
+
+    def serialize(self) -> bytes:
+        return BundledAOTAutogradSerializableCallable.serialize_compile_artifacts(
+            self.inner_fn
+        )
+
+    @staticmethod
+    def deserialize(result_bytes: bytes) -> AOTCompiledArtifact:
+        deserialized = (
+            BundledAOTAutogradSerializableCallable.deserialize_compile_artifacts(
+                result_bytes
+            )
+        )
+        assert isinstance(deserialized, BundledAOTAutogradSerializableCallable)
+        return AOTCompiledArtifact.from_bundled_callable(deserialized)
+
+    @staticmethod
+    def load(
+        *, path: str, format: Literal["binary", "unpacked"] = "binary"
+    ) -> CompiledArtifact:
+        if format == "unpacked":
+            raise RuntimeError(
+                "AOTCompiledArtifact does not support unpacked format yet"
+            )
+        with open(path, "rb") as file:
+            from torch.utils._appending_byte_serializer import BytesReader
+
+            from .codecache import torch_key
+
+            result_bytes = file.read()
+            reader = BytesReader(result_bytes)
+            header = reader.read_bytes()
+            assert header == AOTCompiledArtifact.AOT_HEADER
+            assert reader.read_bytes() == torch_key()
+            artifact = reader.read_bytes()
+            assert reader.is_finished()
+            return AOTCompiledArtifact.deserialize(artifact)
+
+
+def _resolve_ignore_shape_env(dynamic_shapes: Any):
+    # tells compile_fx to ignore the shape_envs on the ambient context
+    # and the graph_module.
+    return dynamic_shapes == "from_example_inputs"
+
+
+def _resolve_fake_mode(gm: GraphModule, dynamic_shapes: Any) -> FakeTensorMode:
+    if dynamic_shapes == "from_example_inputs":
+        return FakeTensorMode(shape_env=ShapeEnv())
+    elif dynamic_shapes == "from_tracing_context":
+        # Reuse fake_mode from the TracingContext.
+        # NB: The TracingContext only exists if we're currently in a torch.compile backend.
+        context = torch._guards.TracingContext.get()
+        assert context.fake_mode is not None
+        return context.fake_mode
+    elif dynamic_shapes == "from_graph":
+        # Strategy: find a FakeTensor in the graph output, grab its FakeTensorMode.
+        # The graph passed to standalone_compile must be an Inductor-approved graph,
+        # which means that there is at least one Tensor output and the output node
+        # contains a flat list of Tensors.
+        last_node = next(iter(reversed(gm.graph.nodes)))
+        assert last_node.op == "output"
+        assert len(last_node.args) == 1
+
+        # If gm came from Dynamo, then last_node.args[0] is always a list,
+        # even in single-Tensor returns.
+        #
+        # It's possible to get into a situation where last_node.args[0]
+        # is a Node (and not a list!). This happens if you call split_module
+        # on the graph. We allow for this case since it is common.
+        nodes = (
+            [last_node.args[0]]
+            if isinstance(last_node.args[0], torch.fx.Node)
+            else last_node.args[0]
+        )
+        for node in nodes:
+            if "example_value" in node.meta:
+                maybe_tensor = node.meta["example_value"]
+                if isinstance(maybe_tensor, torch._subclasses.fake_tensor.FakeTensor):
+                    return maybe_tensor.fake_mode
+
+        return FakeTensorMode(shape_env=ShapeEnv())
+    else:
+        raise ValueError(
+            f"standalone_compile got unsupported `dynamic_shapes` value: dynamic_shapes={dynamic_shapes}."
+        )
+
+
+@contextlib.contextmanager
+def _standalone_context(gm: GraphModule, dynamic_shapes: Any, aot: bool):
+    from torch.compiler._cache import CacheArtifactManager
+
+    fake_mode = _resolve_fake_mode(gm, dynamic_shapes)
+    tracing_context = torch._guards.TracingContext(fake_mode)
+    with (
+        torch._guards.tracing(tracing_context),
+        CacheArtifactManager.with_fresh_cache(),
+        config.patch("triton.autotune_at_compile_time", True),
+        torch._functorch.config.patch("bundled_autograd_cache", aot),
+    ):
+        yield
+
+
+def standalone_compile(
+    gm: GraphModule,
+    example_inputs: Sequence[InputType],
+    *,
+    dynamic_shapes: Any,
+    options: Any,
+    aot: bool = False,  # AOT mode, which uses BundledAOTAutogradCache
+    donate_graph_module: bool = False,
+) -> CompiledArtifact:
+    """
+    Implementation of torch.inductor.standalone_compile
+    """
+    from .compile_fx import compile_fx
+
+    ignore_shape_env = _resolve_ignore_shape_env(dynamic_shapes)
+    with _standalone_context(gm, dynamic_shapes, aot):
+        # compile_fx takes ownership of gm and may mutate it on cache miss.
+        if not donate_graph_module:
+            gm = copy.deepcopy(gm)
+        compiled_fn = compile_fx(
+            gm, example_inputs, ignore_shape_env=ignore_shape_env, **options
+        )
+        assert callable(compiled_fn)
+        if aot:
+            if not hasattr(compiled_fn, "serialize"):
+                raise RuntimeError(
+                    "Compiled function should have serialize method when aot=True"
+                )
+            return AOTCompiledArtifact(compiled_fn)
+        artifacts = torch.compiler.save_cache_artifacts()
+        if artifacts is None:
+            log.warning(
+                "standalone_compile artifact generation failed, cannot save. "
+                "Run with TORCH_LOGS=+torch._inductor.codecache to identify the problem"
+            )
+
+    return CacheCompiledArtifact(compiled_fn, artifacts)
+
+
+def autograd_cache_key(
+    graph,
+    example_inputs,
+    dynamic_shapes: Any,
+    aot: bool = False,  # AOT mode, which uses BundledAOTAutogradCache
+):
+    from . import compile_fx
+
+    ignore_shape_env = _resolve_ignore_shape_env(dynamic_shapes)
+    with _standalone_context(graph, dynamic_shapes, aot):
+        return compile_fx.autograd_cache_key(
+            graph,
+            example_inputs,
+            ignore_shape_env=ignore_shape_env,
+        )

outputs/audit_venv/lib/python3.11/site-packages/torch/_inductor/stream_constants.py

@@ -0,0 +1,5 @@
+"""Constants for CUDA stream naming in Inductor codegen."""
+
+DEFAULT_STREAM: str = "default_stream"
+DEFAULT_STREAM_IDX: int = 0
+STREAM_NAME_TEMPLATE: str = "stream{stream_idx:d}"

outputs/audit_venv/lib/python3.11/site-packages/torch/_inductor/stream_utils.py

@@ -0,0 +1,39 @@
+"""Stream utilities for Inductor codegen."""
+
+from __future__ import annotations
+
+import functools
+
+from torch._inductor.stream_constants import (
+    DEFAULT_STREAM,
+    DEFAULT_STREAM_IDX,
+    STREAM_NAME_TEMPLATE,
+)
+
+
+__all__ = [
+    "DEFAULT_STREAM",
+    "DEFAULT_STREAM_IDX",
+    "STREAM_NAME_TEMPLATE",
+    "get_raw_stream_name",
+    "get_stream_name",
+]
+
+
+@functools.lru_cache
+def get_stream_name(stream_idx: int) -> str:
+    """Generate CUDA Stream name from stream index number.
+
+    Args:
+        stream_idx: Non-negative index number. 0 refers to the default stream, others refer to side
+            streams.
+    """
+    if stream_idx == 0:
+        return DEFAULT_STREAM
+    else:
+        return STREAM_NAME_TEMPLATE.format(stream_idx=stream_idx)
+
+
+def get_raw_stream_name(device_idx: int) -> str:
+    """Generate variable name for a raw stream handle on the given device."""
+    return f"raw_stream{device_idx}"