Diff Coverage

    builds.
    """
    try:
        importlib.import_module("torch.fx.passes.regional_inductor")
    except ImportError:
        return False
    return True


def _get_fake_mode_from_gm(gm: fx.GraphModule) -> Optional[Any]:

        if node.op == "placeholder":
            val = node.meta.get("val")
            if isinstance(val, FakeTensor):
                return val.fake_mode
    return None


def _placeholder_fake_args(gm: fx.GraphModule) -> list:
    """Return the placeholder values in order, requiring fake ``meta["val"]``."""

    args = []
    for node in gm.graph.find_nodes(op="placeholder"):
        val = node.meta.get("val")
        if val is None:
            raise ValueError(
                f"InductorPass requires a fake meta['val'] on placeholder "
                f"'{node.name}' to size compiled regions; the graph must be "
                f"make_fx-traced (PP stage subgraphs copy their placeholders' "
                f"meta from the joint graph)"

    Side-effectful free nodes return ``None`` and legitimately carry no val.
    """
    fake_mode = _get_fake_mode_from_gm(gm)
    if fake_mode is None:
        raise ValueError(
            "InductorPass requires a make_fx-traced graph whose placeholders "
            "carry fake meta['val'] tensors"
        )
    fake_args = _placeholder_fake_args(gm)

        if isinstance(node, fx.Node)
    ):
        return None, graph_module

    flat_nodes: list = []

    def build(node: fx.Node) -> tuple:
        """Emit ``getitem`` leaves for a nested value, else record the leaf."""
        val = node.meta.get("val")
        if isinstance(val, (tuple, list)) and val:
            children = []
            for index, element in enumerate(val):
                sub = graph.call_function(operator.getitem, (node, index))
                sub.meta["val"] = element
                children.append(build(sub))
            kind = "tuple" if isinstance(val, tuple) else "list"
            return kind, tuple(children)
        flat_nodes.append(node)
        return "leaf", len(flat_nodes) - 1

    with graph.inserting_before(output):
        spec = tuple(build(node) for node in returned)
    output.args = (tuple(flat_nodes),)
    graph.lint()
    graph_module.recompile()
    return spec, graph_module


def _rebuild_nested_output(spec: tuple, flat: Sequence[Any]) -> list:
    """Rebuild the original output structure from ``spec`` and flat leaves."""
    leaves = iter(flat)
    rebuilt = []
    for kind, payload in spec:
        if kind == "leaf":
            rebuilt.append(next(leaves))
        else:
            children = _rebuild_nested_output(payload, leaves)
            rebuilt.append(tuple(children) if kind == "tuple" else children)
    return rebuilt


class _NestedOutputWrapper(nn.Module):
    """Thin wrapper restoring a compiled graph's nested output structure."""

        spec: tuple,
        meta: Optional[Dict[str, Any]] = None,
    ) -> None:
        """Store the compiled inner module and its output marker tree."""
        super().__init__()
        self.inner = inner
        self._output_spec = spec
        # Mirrors GraphModule.meta so callers can introspect compilation mode.
        self.meta = dict(meta or {})

    def forward(self, *args: Any, **kwargs: Any) -> tuple:
        """Run the flat compiled graph and rebuild the original structure."""
        flat = self.inner(*args, **kwargs)
        return tuple(_rebuild_nested_output(self._output_spec, flat))


def _tag_node(node: fx.Node, inductor_configs: Optional[Dict[str, Any]]) -> None:
    """Annotate ``node`` for regional_inductor with the user's Inductor config."""

    @property
    def name(self) -> str:
        """Pass name."""
        return "inductor_compilation"

    def run(
        self,
        graph_module: fx.GraphModule,

                    "JointGraph.example_inputs) to size compiled regions"
                )
            if _regional_inductor_available():
                self._compile_graph(graph_module, pass_config, example_inputs)
                compiled = True
            else:
                self._warn_unavailable(mode)

        graph_module.meta[_INDUCTOR_COMPILED_META] = (

                )
            schedules.append(submodule)

        if not _regional_inductor_available():
            self._warn_unavailable(pass_config.inductor_mode)
            return False

        for schedule in schedules:
            schedule.map_stage_graphs(
                lambda stage_gm: self._compile_graph(

                lambda stage_gm: self._compile_graph(
                    stage_gm, pass_config, example_inputs=None
                )
            )
        _LOG.info(
            "InductorPass: compiled %d PP schedule(s) in %r mode",
            len(schedules),
            pass_config.inductor_mode,
        )
        return True

    def _compile_graph(
        self,
        graph_module: fx.GraphModule,

        ):
            # ``regional_inductor`` consumes these on torch releases whose
            # signature expects example inputs; newer releases size regions
            # from the nodes' fake ``meta["val"]`` instead.
            regional_inductor(graph_module, example_inputs)

        # regional_inductor compiles regions with the boxed calling
        # convention (one mutable args list); reset to the default positional
        # convention so run_traced_graph's graph_module(*flat_inputs) (and the

        # convention so run_traced_graph's graph_module(*flat_inputs) (and the
        # PP schedule's fwd_gm(*args) / bwd_gm(*args)) keeps working. Setting
        # the default unconditionally is a no-op when the region compiler
        # already used positional codegen and avoids peeking at private state.
        graph_module.graph.set_codegen(CodeGen())
        graph_module.recompile()
        graph_module.meta[_INDUCTOR_COMPILED_META] = mode
        if output_spec is not None:
            # Rebuild the stage's original (possibly nested) output structure
            # for the schedule / backward subgraph; only the compiled graph is
            # flat.
            return _NestedOutputWrapper(graph_module, output_spec, graph_module.meta)
        return graph_module

    @staticmethod
    def _tag_full(
        graph_module: fx.GraphModule, inductor_configs: Optional[Dict[str, Any]]

        them through this hook and writes the returned module back (in-place
        compilers may return the same object).
        """
        self.fwd_gm = fn(self.fwd_gm)
        self.bwd_gm = fn(self.bwd_gm)

    # ------------------------------------------------------------------
    # Input / microbatch resolution
    # ------------------------------------------------------------------

        self._pending_sends: Dict[Any, List[Tuple[Any, torch.Tensor]]] = {}

    def map_stage_graphs(self, fn: Callable[[nn.Module], nn.Module]) -> None:
        """Apply ``fn`` to every virtual stage's fwd and bwd subgraph."""
        self.fwd_gms = nn.ModuleList([fn(gm) for gm in self.fwd_gms])
        self.bwd_gms = nn.ModuleList([fn(gm) for gm in self.bwd_gms])

    def _resolve_num_microbatches(self, user_inputs: Sequence[Any]) -> int:
        """Derive the microbatch count from the first tensor model input."""
        batch_size = next(