procfunc 0.30.0__py3-none-any.whl

procfunc/compute_graph/proxy.py

@@ -0,0 +1,173 @@
+"""General-purpose Proxy wrapper for Node with all dunders."""
+
+import operator
+from dataclasses import dataclass
+from typing import Generic, TypeVar
+
+from .node import FunctionCallNode, GetAttributeNode, MethodCallNode, Node
+
+T = TypeVar("T")
+
+
+@dataclass
+class Proxy(Generic[T]):
+    """General-purpose wrapper for Node that provides all dunder methods."""
+
+    node: Node
+
+    def __repr__(self):
+        return f"Proxy({self.node!r})"
+
+    def __getattr__(self, attr: str) -> "AttributeProxy":
+        node = GetAttributeNode(source=self.node, attribute_name=attr)
+        return AttributeProxy(node)
+
+    def __call__(self, *args, **kwargs) -> "Proxy":
+        raise NotImplementedError("Proxy.__call__ is not implemented")
+
+    def __len__(self) -> int:
+        raise ValueError(
+            "Tracing does not allow __len__ since real values are not evaluated"
+        )
+
+    def __iter__(self):
+        raise ValueError(
+            "Proxy does not support __iter__. Use explicit indexing instead."
+        )
+
+    def __getitem__(self, idx) -> "Proxy":
+        idx_node = idx.node if isinstance(idx, Proxy) else idx
+        getitem_node = FunctionCallNode(
+            func=operator.getitem,
+            args=(self.node, idx_node),
+            kwargs={},
+        )
+        return Proxy(getitem_node)
+
+    def __bool__(self):
+        raise ValueError(
+            "Base Proxy does not allow __bool__ during tracing since real values are unknown"
+        )
+
+
+NODE_DUNDER_METHODS = {
+    "__add__": operator.add,
+    "__sub__": operator.sub,
+    "__mul__": operator.mul,
+    "__truediv__": operator.truediv,
+    "__floordiv__": operator.floordiv,
+    "__mod__": operator.mod,
+    "__pow__": operator.pow,
+    "__lshift__": operator.lshift,
+    "__rshift__": operator.rshift,
+    "__and__": operator.and_,
+    "__xor__": operator.xor,
+    "__or__": operator.or_,
+    "__neg__": operator.neg,
+    "__pos__": operator.pos,
+    "__abs__": operator.abs,
+    "__invert__": operator.invert,
+    "__eq__": operator.eq,
+    "__ne__": operator.ne,
+    "__lt__": operator.lt,
+    "__le__": operator.le,
+    "__gt__": operator.gt,
+    "__ge__": operator.ge,
+}
+
+NODE_REFLECTABLE_METHODS = [
+    "add",
+    "sub",
+    "mul",
+    "floordiv",
+    "truediv",
+    "div",
+    "mod",
+    "pow",
+    "lshift",
+    "rshift",
+    "and_",
+    "or_",
+    "xor",
+    "getitem",
+    "matmul",
+]
+
+
+def _add_proxy_operator(cls, name, operator_func):
+    def proxy_method(self, *args, **kwargs):
+        # Convert any Proxy args to their underlying nodes
+        node_args = tuple(arg.node if isinstance(arg, Proxy) else arg for arg in args)
+        node_kwargs = {
+            k: v.node if isinstance(v, Proxy) else v for k, v in kwargs.items()
+        }
+        node = FunctionCallNode(
+            func=operator_func,
+            args=(self.node, *node_args),
+            kwargs=node_kwargs,
+        )
+        return Proxy(node)
+
+    setattr(cls, name, proxy_method)
+
+
+def _add_proxy_reflection(cls, name: str):
+    # __rmul__(self, rhs) means rhs * self — use the same operator but swap arg order
+    fwd_dunder = f"__{name.rstrip('_')}__"
+    operator_func = NODE_DUNDER_METHODS.get(fwd_dunder)
+    if operator_func is None:
+        return  # no matching forward op, skip
+
+    def proxy_method(self, rhs):
+        rhs_node = rhs.node if isinstance(rhs, Proxy) else rhs
+        node = FunctionCallNode(
+            func=operator_func,
+            args=(rhs_node, self.node),
+            kwargs={},
+        )
+        return Proxy(node)
+
+    setattr(cls, f"__r{name}__", proxy_method)
+
+
+# Add all dunder methods to Proxy
+for name, operator_func in NODE_DUNDER_METHODS.items():
+    _add_proxy_operator(Proxy, name, operator_func)
+
+# Add reflected methods
+for name in NODE_REFLECTABLE_METHODS:
+    _add_proxy_reflection(Proxy, name)
+
+
+@dataclass
+class AttributeProxy(Proxy):
+    """Special proxy for attribute access that supports peekthrough optimization"""
+
+    def __init__(self, node: Node):
+        super().__init__(node)
+        assert isinstance(node, GetAttributeNode), node
+
+    def __call__(self, *args, **kwargs) -> Proxy:
+        """
+        Someone did func = proxy.xyz, then func(), or equivelantly thats just proxy.xyz().
+        We can convert that to just a single node which is a method call on the obj node.
+
+        torch.fx.symbolic_trace calls this a _peekthrough optimization_
+
+        TODO: this means that `self` is very often dropped from the graph. need to account for this if we check for drops.
+        """
+        assert isinstance(self.node, GetAttributeNode), self.node
+
+        # Convert any Proxy args to their underlying nodes
+        node_args = tuple(arg.node if isinstance(arg, Proxy) else arg for arg in args)
+        node_kwargs = {
+            k: v.node if isinstance(v, Proxy) else v for k, v in kwargs.items()
+        }
+        # self.node.args[0] is the source node that we're calling the method on
+        call_node = MethodCallNode(
+            callee=self.node.args[0],
+            method_name=self.node.attribute_name,
+            args=node_args,
+            kwargs=node_kwargs,
+        )
+        return Proxy(call_node)

procfunc/compute_graph/util.py

@@ -0,0 +1,282 @@
+import copy
+import logging
+from collections import defaultdict, deque
+from typing import Any, Callable, Generator, Literal, TypeVar
+
+from procfunc.util import pytree
+
+from .compute_graph import ComputeGraph
+from .node import Node, SubgraphCallNode
+
+logger = logging.getLogger(__name__)
+
+
+T = TypeVar("T")
+
+
+class LiteralConstant:
+    def __init__(self, value: Any):
+        self.value = value
+
+    def __repr__(self) -> Any:
+        return self.value
+
+
+def traverse_breadth_first(
+    graph: ComputeGraph,
+    yield_parent: bool = False,
+    yield_name: bool = False,
+    yield_consts: bool = False,
+) -> Generator[Any, None, None]:
+    """
+    Traverse all nodes in the compute graph.
+
+    Args:
+        graph: The compute graph to traverse
+        yield_parent: If True, yield (parent, child), with output nodes having parent=None
+        yield_name: If True, yield (name, child) or (name, parent, child) if yield_parent is also True
+        yield_consts: If True, yield child arguments of nodes even if they are not Nodes
+    """
+
+    visited = set()
+    frontier = deque((None, name, node) for name, node in graph.outputs.items())
+    # logger.debug(f"{traverse_breadth_first.__name__} {graph.name} {len(frontier)=}")
+
+    def res(parent, name, child):
+        res = (child,)
+        if yield_parent:
+            res = (parent,) + res
+        if yield_name:
+            res = (name,) + res
+
+        return res[0] if len(res) == 1 else tuple(res)
+
+    while len(frontier) > 0:
+        parent, name, node = frontier.popleft()
+
+        if yield_consts and not isinstance(node, Node):
+            yield res(parent, name, node)
+            continue
+
+        if node is None:
+            continue
+
+        if id(node) in visited:
+            continue
+        visited.add(id(node))
+
+        yield res(parent, name, node)
+
+        children = list(pytree.PyTree(node.args).items()) + list(
+            pytree.PyTree(node.kwargs).items()
+        )
+        for key, arg in children:
+            if not yield_consts and not isinstance(arg, Node):
+                continue
+            if id(arg) in visited:
+                continue
+            frontier.append((node, key, arg))
+
+
+def _traverse_depth_first_node(
+    node: Node,
+    visited: set[int],
+    parent: Node | None,
+    name: str,
+    order: Literal["preorder", "postorder"],
+    yield_parent: bool,
+    yield_name: bool,
+    yield_consts: bool,
+) -> Generator[Any, None, None]:
+    def res(parent, name, child):
+        res = (child,)
+        if yield_parent:
+            res = (parent,) + res
+        if yield_name:
+            res = (name,) + res
+        return res[0] if len(res) == 1 else tuple(res)
+
+    assert isinstance(node, Node), node
+    if id(node) in visited:
+        return
+    visited.add(id(node))
+
+    if order == "preorder":
+        yield res(parent, name, node)
+
+    children = list(pytree.PyTree(node.args).items()) + list(
+        pytree.PyTree(node.kwargs).items()
+    )
+
+    for key, arg in children:
+        if not isinstance(arg, Node):
+            if yield_consts:
+                yield res(node, key, arg)
+            continue
+        yield from _traverse_depth_first_node(
+            node=arg,
+            visited=visited,
+            parent=node,
+            name=key,
+            order=order,
+            yield_parent=yield_parent,
+            yield_name=yield_name,
+            yield_consts=yield_consts,
+        )
+
+    if order == "postorder":
+        yield res(parent, name, node)
+
+
+def traverse_depth_first_node(
+    node: Node,
+    yield_consts: bool = False,
+    order: Literal["preorder", "postorder"] = "postorder",
+) -> Generator[Any, None, None]:
+    return _traverse_depth_first_node(
+        node=node,
+        visited=set(),
+        parent=None,
+        name="",
+        order=order,
+        yield_parent=False,
+        yield_name=False,
+        yield_consts=yield_consts,
+    )
+
+
+def traverse_depth_first(
+    graph: ComputeGraph,
+    yield_parent: bool = False,
+    yield_name: bool = False,
+    yield_consts: bool = False,
+    order: Literal["preorder", "postorder"] = "postorder",
+) -> Generator[Any, None, None]:
+    visited = set()
+    for name, node in graph.outputs.items():
+        if node is None:
+            continue
+        yield from _traverse_depth_first_node(
+            node, visited, None, name, order, yield_parent, yield_name, yield_consts
+        )
+
+
+def traverse_nested_graphs(
+    graph: ComputeGraph,
+    yield_call_nodes: bool = False,
+) -> Generator[tuple[Node | None, ComputeGraph], None, None]:
+    visited = set()
+    frontier = deque([(None, graph)])
+
+    while len(frontier) > 0:
+        node, graph = frontier.popleft()
+
+        if id(graph) in visited:
+            continue
+        visited.add(id(graph))
+
+        if yield_call_nodes:
+            yield node, graph
+        else:
+            yield graph
+
+        frontier.extend(
+            (node, node.subgraph)
+            for node in traverse_depth_first(graph)
+            if isinstance(node, SubgraphCallNode)
+        )
+
+
+def usages_per_node(
+    graph: ComputeGraph,
+) -> dict[int, list[Node]]:
+    usages = defaultdict(list)
+    for node in traverse_depth_first(graph):
+        argtree = pytree.PyTree((node.args, node.kwargs))
+        for arg in argtree.values():
+            if isinstance(arg, Node):
+                usages[id(arg)].append(node)
+    return dict(usages)
+
+
+def graph_nodes_equal(graph1: ComputeGraph, graph2: ComputeGraph) -> bool:
+    nodes1 = list(traverse_depth_first(graph1))
+    nodes2 = list(traverse_depth_first(graph2))
+    if len(nodes1) != len(nodes2):
+        return False
+    for node1, node2 in zip(nodes1, nodes2):
+        if type(node1) is not type(node2):
+            return False
+        if isinstance(node1, SubgraphCallNode):
+            if not graph_nodes_equal(node1.subgraph, node2.subgraph):
+                return False
+        elif node1.args != node2.args or node1.kwargs != node2.kwargs:
+            return False
+    return True
+
+
+def transform_nodetree(
+    root: Node,
+    transform_fn: Callable[[Node], Any],
+    memo: dict[int, Node] = {},
+):
+    raise NotImplementedError("Not implemented")
+
+    new_root = transform_fn(root)
+
+    for parent, parent_key, node in traverse_breadth_first(root, parent_child=True):
+        if parent is None:
+            continue
+        elif parent is root:
+            parent = new_root
+
+        new_node = transform_fn(node)
+        if new_node is None:
+            raise ValueError(
+                f"Transform function {transform_fn.__name__} returned None for node {node.name}"
+            )
+        if isinstance(parent_key, int):
+            args_list = list(parent.args)
+            args_list[parent_key] = new_node
+            parent.args = tuple(args_list)
+        else:
+            parent.kwargs[parent_key] = new_node
+
+    return new_root
+
+
+def transform_compute_graph(
+    compute_graph: ComputeGraph,
+    transform_fn: Callable[[Node], Any],
+    graph_name: str | None = None,
+):
+    raise NotImplementedError("Not implemented")
+
+    id_map: dict[int, Node] = {}
+
+    def wrapper(node: Node) -> Node:
+        res = transform_fn(node)
+        id_map[id(node)] = res
+        return res
+
+    memo = {}
+
+    new_output_values = [
+        transform_nodetree(v, wrapper, memo) for v in compute_graph.outputs.values()
+    ]
+    new_output = pytree.PyTree.from_values(new_output_values, compute_graph.output.spec)
+
+    new_metadata = copy.copy(compute_graph.metadata)
+    if "operations" not in new_metadata:
+        new_metadata["operations"] = []
+    op = (transform_compute_graph, {"transform_fn": transform_fn, "id_map": id_map})
+    new_metadata["operations"].append(op)
+
+    new_inputs = compute_graph.inputs.map(lambda v: id_map[id(v)])
+
+    return ComputeGraph(
+        inputs=new_inputs,
+        outputs=new_output,
+        name=compute_graph.name + "_transformed",
+        metadata=new_metadata,
+    )

procfunc/context.py

@@ -0,0 +1,115 @@
+import copy
+import multiprocessing
+import os
+from contextlib import contextmanager
+from dataclasses import asdict, dataclass
+from typing import Any, Literal
+
+
+@dataclass
+class ProcfuncContext:
+    """Global context for Procfunc configuration and system information."""
+
+    num_cpu_cores: int
+    current_trace_level: int | None  # compared against to TraceLevel int values
+
+    warn_mode_avoid_normal_bump: Literal["ignore", "warn", "throw"]
+    """
+    Set to 'throw' for infinigen best practices. Using a normal map / bump map / BSDF vector input _can_ be useful,
+    but its always better to do it via the Displacement output of the shader (so that it at least can also be done as displacement)
+    """
+
+    warn_mode_avoid_implicit_vector: Literal["ignore", "warn", "throw"]
+    """
+    Set to 'throw' for default infinigen usage - we force the user to specify exactly what vector to sample
+    This prevents materials from being incorrect when on a moving object, or ignoring uv coordinates
+    """
+
+    warn_mode_avoid_io_nodes: Literal["ignore", "warn", "throw"]
+    """
+    Set to 'throw' for infinigen, prevents floating Value Vector Color nodes, or strange output nodes like AOV in shader
+    Intent is to force nodegroups to be more functional - all inputs and outputs go through nodegroup interface
+    """
+
+    warn_mode_empty_geonodes: Literal["ignore", "warn", "throw"]
+    """
+    Controls behavior when a geometry node graph produces no mesh geometry (e.g. unconnected inputs).
+    'ignore' silently returns an empty mesh, 'warn' logs a warning, 'throw' raises an error.
+    """
+
+    def __post_init__(self):
+        """Initialize computed fields after dataclass creation."""
+        if self.num_cpu_cores <= 0:
+            self.num_cpu_cores = multiprocessing.cpu_count()
+
+    def set_strict(self):
+        """Set all warning modes to 'throw'"""
+        self.warn_mode_avoid_normal_bump = "throw"
+        self.warn_mode_avoid_implicit_vector = "throw"
+        self.warn_mode_avoid_io_nodes = "throw"
+        self.warn_mode_empty_geonodes = "throw"
+
+    def set_warn(self):
+        """Set all warning modes to 'warn'"""
+        self.warn_mode_avoid_normal_bump = "warn"
+        self.warn_mode_avoid_implicit_vector = "warn"
+        self.warn_mode_avoid_io_nodes = "warn"
+        self.warn_mode_empty_geonodes = "warn"
+
+
+# Global context instance
+
+warn_modes = ["ignore", "warn", "throw"]
+
+_warn_mode_avoid_normal_bump = os.environ.get(
+    "PROCFUNC_WARN_MODE_AVOID_NORMAL_BUMP", "ignore"
+)
+assert _warn_mode_avoid_normal_bump in warn_modes
+
+warn_mode_avoid_implicit_vector = os.environ.get(
+    "PROCFUNC_WARN_MODE_AVOID_IMPLICIT_VECTOR", "ignore"
+)
+assert warn_mode_avoid_implicit_vector in warn_modes
+
+warn_mode_avoid_io_nodes = os.environ.get("PROCFUNC_WARN_MODE_AVOID_IO_NODES", "ignore")
+assert warn_mode_avoid_io_nodes in warn_modes
+
+_warn_mode_empty_geonodes = os.environ.get("PROCFUNC_WARN_MODE_EMPTY_GEONODES", "warn")
+assert _warn_mode_empty_geonodes in warn_modes
+
+globals = ProcfuncContext(
+    num_cpu_cores=int(os.environ.get("PROCFUNC_NUM_CPU_CORES", 0)),
+    warn_mode_avoid_normal_bump=_warn_mode_avoid_normal_bump,  # type: ignore[invalid-assignment]
+    warn_mode_avoid_implicit_vector=warn_mode_avoid_implicit_vector,  # type: ignore[invalid-assignment]
+    warn_mode_avoid_io_nodes=warn_mode_avoid_io_nodes,  # type: ignore[invalid-assignment]
+    warn_mode_empty_geonodes=_warn_mode_empty_geonodes,  # type: ignore[invalid-assignment]
+    current_trace_level=None,
+)
+
+
+@contextmanager
+def override_globals(
+    new_context: ProcfuncContext | None = None,
+    **overrides: Any,
+):
+    """
+    Override the context for a block of code
+
+    Args:
+        new_context: If provided, will override the entire context with this new context
+        overrides: If provided, will override specific keys with these values
+    """
+    orig = copy.deepcopy(globals)
+
+    if new_context is not None:
+        for key, value in asdict(new_context).items():
+            setattr(globals, key, value)
+
+    if overrides is not None:
+        for key, value in overrides.items():
+            setattr(globals, key, value)
+
+    yield
+
+    for key, value in asdict(orig).items():
+        setattr(globals, key, value)