AI: optimalisation of script input node result
This commit is contained in:
mcDandy
@@ -0,0 +1,228 @@
|
||||
from antlr4.tree.Tree import TerminalNode
|
||||
|
||||
from .antlr_router import get_antlr_modules
|
||||
|
||||
_, P, _ = get_antlr_modules()
|
||||
|
||||
# builtins that print, touch the stack or change their argument in place; calls to them are never removed
|
||||
_EFFECT_TOKENS = {
|
||||
P.PRNT, P.PRINT_SHAPE, P.PUSH, P.POP, P.CLEAR, P.HAS, P.GET,
|
||||
P.ADD_KEY, P.REMOVE_KEY, P.ENABLE_TRACKING, P.DISABLE_TRACKING, P.RESET_TRACKING,
|
||||
}
|
||||
|
||||
|
||||
def remove_dead_code(tree):
|
||||
"""Removes statements whose result is never read and that have no side effects. Changes and returns tree."""
|
||||
_DeadCode(tree).run()
|
||||
return tree
|
||||
|
||||
|
||||
class _DeadCode:
|
||||
def __init__(self, tree):
|
||||
self.tree = tree
|
||||
self.defs = {}
|
||||
self.lambdas = []
|
||||
self.calls = []
|
||||
assignments = []
|
||||
# which lambdas a variable can hold, and variables that may hold one we cannot follow
|
||||
self.lam = {}
|
||||
self.unknown = set()
|
||||
todo = [tree]
|
||||
while todo:
|
||||
n = todo.pop()
|
||||
if isinstance(n, TerminalNode):
|
||||
continue
|
||||
if isinstance(n, P.VarDefContext):
|
||||
if n.EQUEALS() and not n.LBRACKET():
|
||||
assignments.append((n.VARIABLE().getText(), n.expr(0)))
|
||||
elif isinstance(n, P.ForStmtContext):
|
||||
self.unknown.add(n.VARIABLE().getText())
|
||||
elif isinstance(n, P.FunctionDefContext):
|
||||
self.defs.setdefault(n.VARIABLE().getText(), []).append((_params(n), n.block() or n.expr()))
|
||||
elif isinstance(n, P.LambdaExpContext):
|
||||
self.lambdas.append(n)
|
||||
elif isinstance(n, P.CallExpContext):
|
||||
self.calls.append(n)
|
||||
todo.extend(n.getChildren())
|
||||
|
||||
# User functions see the caller's variables, so a call reads everything its body reads.
|
||||
# Lambdas only see the variables captured when they were created, so calling one reads nothing extra.
|
||||
self.fn_reads = {name: set() for name in self.defs}
|
||||
self.fn_effect = {name: False for name in self.defs}
|
||||
self.lambda_effect = {lam: False for lam in self.lambdas}
|
||||
changed = True
|
||||
while changed:
|
||||
changed = False
|
||||
for name, expr in assignments:
|
||||
changed |= self.flow(name, self.value(expr))
|
||||
for call in self.calls:
|
||||
lams, fns, _ = self.resolve(call.VARIABLE().getText())
|
||||
args = call.exprList().expr() if call.exprList() else []
|
||||
targets = [_params(lam) for lam in lams] + [params for fn in fns for params, _ in self.defs[fn]]
|
||||
for params in targets:
|
||||
for param, arg in zip(params, args):
|
||||
changed |= self.flow(param, self.value(arg))
|
||||
for lam in self.lambdas:
|
||||
if not self.lambda_effect[lam] and self.scan(lam.block() or lam.expr())[1]:
|
||||
self.lambda_effect[lam] = changed = True
|
||||
for name, defs in self.defs.items():
|
||||
reads, effect = set(), False
|
||||
for params, body in defs:
|
||||
r, e = self.scan(body)
|
||||
reads |= r - set(params)
|
||||
effect |= e
|
||||
if reads != self.fn_reads[name] or effect != self.fn_effect[name]:
|
||||
self.fn_reads[name], self.fn_effect[name] = reads, effect
|
||||
changed = True
|
||||
|
||||
def flow(self, name, value):
|
||||
lams, unknown = value
|
||||
changed = not lams <= self.lam.get(name, set()) or (unknown and name not in self.unknown)
|
||||
self.lam.setdefault(name, set()).update(lams)
|
||||
if unknown:
|
||||
self.unknown.add(name)
|
||||
return changed
|
||||
|
||||
def value(self, node):
|
||||
"""Lambdas an expression can evaluate to, and whether it may be one we cannot follow."""
|
||||
while True:
|
||||
if isinstance(node, P.ParenExpContext):
|
||||
node = node.expr()
|
||||
elif node.getChildCount() == 1 and not isinstance(node.getChild(0), TerminalNode):
|
||||
node = node.getChild(0)
|
||||
else:
|
||||
break
|
||||
if isinstance(node, P.LambdaExpContext):
|
||||
return {node}, False
|
||||
if isinstance(node, P.VariableExpContext):
|
||||
name = node.getText()
|
||||
return set(self.lam.get(name, ())), name in self.unknown
|
||||
if isinstance(node, P.TernaryExpContext):
|
||||
(a, ua), (b, ub) = self.value(node.expr(0)), self.value(node.expr(1))
|
||||
return a | b, ua or ub
|
||||
todo = [node]
|
||||
while todo:
|
||||
n = todo.pop()
|
||||
if isinstance(n, TerminalNode):
|
||||
continue
|
||||
if isinstance(n, (P.LambdaExpContext, P.CallExpContext)):
|
||||
return set(), True
|
||||
if isinstance(n, P.VariableExpContext) and (self.lam.get(n.getText()) or n.getText() in self.unknown):
|
||||
return set(), True
|
||||
todo.extend(n.getChildren())
|
||||
return set(), False
|
||||
|
||||
def resolve(self, name):
|
||||
"""Lambdas and user functions a call by this name can reach, and whether it may reach something else."""
|
||||
lams = self.lam.get(name, set())
|
||||
unknown = name in self.unknown
|
||||
fns = [name] if name in self.defs else []
|
||||
return lams, fns, unknown or not (lams or fns)
|
||||
|
||||
def run(self):
|
||||
live, _ = self.scan(self.tree.expr())
|
||||
self.sweep(self.tree, live, None, apply=True, keep_last=False, skip=self.tree.expr())
|
||||
|
||||
def scan(self, node):
|
||||
"""Variables an expression reads, and whether evaluating it has side effects."""
|
||||
reads, effect = set(), False
|
||||
todo = [node]
|
||||
while todo:
|
||||
n = todo.pop()
|
||||
if isinstance(n, TerminalNode):
|
||||
t = n.getSymbol().type
|
||||
if t == P.VARIABLE:
|
||||
reads.add(n.getText())
|
||||
elif t in _EFFECT_TOKENS:
|
||||
effect = True
|
||||
continue
|
||||
if isinstance(n, P.LambdaExpContext):
|
||||
# creating a lambda only captures variables; its body runs when it is called
|
||||
reads |= self.scan(n.block() or n.expr())[0] - set(_params(n))
|
||||
continue
|
||||
if isinstance(n, P.CallExpContext):
|
||||
lams, fns, unknown = self.resolve(n.VARIABLE().getText())
|
||||
effect |= unknown or any(self.lambda_effect[lam] for lam in lams) or any(self.fn_effect[fn] for fn in fns)
|
||||
for fn in fns:
|
||||
reads |= self.fn_reads[fn]
|
||||
todo.extend(n.getChildren())
|
||||
return reads, effect
|
||||
|
||||
def sweep(self, parent, live, loop_live, apply, keep_last, skip=None):
|
||||
"""Walks a statement list backwards, removing dead statements when apply is set; returns live-before."""
|
||||
stmts = [c for c in parent.children if not isinstance(c, TerminalNode) and c is not skip]
|
||||
for i in range(len(stmts) - 1, -1, -1):
|
||||
before, removable = self.stmt(stmts[i], live, loop_live, apply)
|
||||
if removable and not (keep_last and i == len(stmts) - 1):
|
||||
if apply:
|
||||
parent.children.remove(stmts[i])
|
||||
else:
|
||||
live = before
|
||||
return live
|
||||
|
||||
def stmt(self, node, live, loop_live, apply):
|
||||
"""Returns (live-before if kept, whether the statement can be removed)."""
|
||||
if isinstance(node, P.VarDefStmtContext):
|
||||
node = node.varDef()
|
||||
if isinstance(node, P.VarDefContext):
|
||||
name = node.VARIABLE().getText()
|
||||
reads, effect = set(), False
|
||||
for e in node.expr():
|
||||
r, eff = self.scan(e)
|
||||
reads |= r
|
||||
effect |= eff
|
||||
if node.EQUEALS() and not node.LBRACKET():
|
||||
before = (live - {name}) | reads
|
||||
else:
|
||||
before = live | {name} | reads
|
||||
return before, name not in live and not effect
|
||||
if isinstance(node, P.ExprStatementContext):
|
||||
reads, effect = self.scan(node.expr())
|
||||
return live | reads, not effect
|
||||
if isinstance(node, P.IfStatementContext):
|
||||
node = node.ifStmt()
|
||||
reads, _ = self.scan(node.expr())
|
||||
before = reads | self.stmt_kept(node.stmt(0), live, loop_live, apply)
|
||||
if node.stmt(1):
|
||||
return before | self.stmt_kept(node.stmt(1), live, loop_live, apply), False
|
||||
return before | live, False
|
||||
if isinstance(node, P.WhileStatementContext):
|
||||
node = node.whileStmt()
|
||||
reads, _ = self.scan(node.expr())
|
||||
head = self.loop_head(node.stmt(), live | reads, None, apply)
|
||||
return head, False
|
||||
if isinstance(node, P.ForStatementContext):
|
||||
node = node.forStmt()
|
||||
reads, _ = self.scan(node.expr())
|
||||
head = self.loop_head(node.stmt(), live, node.VARIABLE().getText(), apply)
|
||||
return head | reads, False
|
||||
if isinstance(node, P.BlockStatementContext):
|
||||
return self.sweep(node.block(), live, loop_live, apply, keep_last=True), False
|
||||
if isinstance(node, (P.BreakStatementContext, P.ContinueStatementContext)):
|
||||
return set(loop_live), False
|
||||
if isinstance(node, P.ReturnStatementContext):
|
||||
expr = node.returnStmt().expr()
|
||||
return (self.scan(expr)[0] if expr else set()), False
|
||||
# function definitions and anything unexpected stay, and are assumed to read everything still live
|
||||
return live, False
|
||||
|
||||
def stmt_kept(self, node, live, loop_live, apply):
|
||||
return self.stmt(node, live, loop_live, apply)[0]
|
||||
|
||||
def loop_head(self, body, exit_live, loop_var, apply):
|
||||
"""Fixed point of the variables live at the top of a loop; only removes once it is stable."""
|
||||
head = set(exit_live)
|
||||
while True:
|
||||
body_live = self.stmt_kept(body, head, head, apply=False)
|
||||
if loop_var is not None:
|
||||
body_live = body_live - {loop_var}
|
||||
new_head = exit_live | body_live
|
||||
if new_head == head:
|
||||
break
|
||||
head = new_head
|
||||
self.stmt_kept(body, head, head, apply)
|
||||
return head
|
||||
|
||||
|
||||
def _params(node):
|
||||
return [v.getText() for v in node.paramList().VARIABLE()] if node.paramList() else []
|
||||
@@ -1,6 +1,7 @@
|
||||
from comfy_api.latest import io
|
||||
from .helper_functions import parse_expr
|
||||
from .ParseTree import MrmthParseTree
|
||||
from .Parser.dead_code import remove_dead_code
|
||||
|
||||
class ScriptTextInput(io.ComfyNode):
|
||||
|
||||
@@ -21,5 +22,4 @@ class ScriptTextInput(io.ComfyNode):
|
||||
|
||||
@classmethod
|
||||
def execute(cls, script):
|
||||
expr = parse_expr(script)
|
||||
return (expr,)
|
||||
return (remove_dead_code(parse_expr(script)),)
|
||||
|
||||
@@ -1,5 +1,8 @@
|
||||
from antlr4.error.ErrorListener import ErrorListener
|
||||
from antlr4 import InputStream, CommonTokenStream
|
||||
from antlr4.atn.PredictionMode import PredictionMode
|
||||
from antlr4.error.ErrorStrategy import BailErrorStrategy, DefaultErrorStrategy
|
||||
from antlr4.error.Errors import ParseCancellationException
|
||||
|
||||
import torch
|
||||
|
||||
@@ -69,8 +72,20 @@ def parse_expr(expr: str):
|
||||
lexer = MathExprLexer(input_stream)
|
||||
stream = CommonTokenStream(lexer)
|
||||
parser = MathExprParser(stream)
|
||||
parser.addErrorListener(ThrowingErrorListener())
|
||||
return parser.start()
|
||||
# SLL prediction is orders of magnitude faster and gives the same tree for valid input;
|
||||
# only fall back to full LL (with proper error messages) when it bails out.
|
||||
parser._interp.predictionMode = PredictionMode.SLL
|
||||
parser.removeErrorListeners()
|
||||
parser._errHandler = BailErrorStrategy()
|
||||
try:
|
||||
return parser.start()
|
||||
except ParseCancellationException:
|
||||
stream.seek(0)
|
||||
parser.reset()
|
||||
parser._interp.predictionMode = PredictionMode.LL
|
||||
parser._errHandler = DefaultErrorStrategy()
|
||||
parser.addErrorListener(ThrowingErrorListener())
|
||||
return parser.start()
|
||||
|
||||
|
||||
def getIndexTensorAlongDim(tensor, dim):
|
||||
|
||||
@@ -0,0 +1,113 @@
|
||||
import io
|
||||
import contextlib
|
||||
|
||||
import torch
|
||||
|
||||
from more_math.Parser.UnifiedMathVisitor import UnifiedMathVisitor
|
||||
from more_math.Parser.dead_code import remove_dead_code
|
||||
from more_math.helper_functions import parse_expr, checkLazyNew
|
||||
|
||||
|
||||
def run(script, variables=None, optimize=True, stack=None):
|
||||
tree = parse_expr(script)
|
||||
if optimize:
|
||||
remove_dead_code(tree)
|
||||
out = io.StringIO()
|
||||
with contextlib.redirect_stdout(out):
|
||||
result = UnifiedMathVisitor(dict(variables or {}), (1,), state_storage=stack if stack is not None else {}).visit(tree)
|
||||
return result, out.getvalue(), tree.getText()
|
||||
|
||||
|
||||
def same_as_unoptimized(script, variables=None):
|
||||
plain, plain_out, _ = run(script, variables, optimize=False, stack={})
|
||||
result, out, text = run(script, variables, stack={})
|
||||
assert out == plain_out
|
||||
if torch.is_tensor(plain):
|
||||
torch.testing.assert_close(result, plain)
|
||||
else:
|
||||
assert result == plain
|
||||
return text
|
||||
|
||||
|
||||
def test_removes_unused_assignments_and_expression_statements():
|
||||
text = same_as_unoptimized("a = V0 * 2; b = V0 + 1; V0 * 3; b", {"V0": torch.ones(4)})
|
||||
assert "a=" not in text and "V0*3;" not in text and "b=V0+1" in text
|
||||
|
||||
|
||||
def test_keeps_side_effects():
|
||||
stack = {1: [7]}
|
||||
result, out, text = run("x = print(1); stack_push(1, 2); y = stack_pop(1) + stack_get(1); pshp(3); 5", stack=stack)
|
||||
assert "x=print(1)" in text and "stack_push(1,2)" in text and "y=" in text and "pshp(3)" in text
|
||||
assert result == 5 and stack[1] == [7] and "1" in out
|
||||
|
||||
|
||||
def test_keeps_variables_read_by_called_functions():
|
||||
text = same_as_unoptimized("p = V0 * 2; q = V0 * 9; f() -> p + 1; f()", {"V0": torch.ones(2)})
|
||||
assert "p=V0*2" in text and "q=" not in text
|
||||
|
||||
|
||||
def test_function_with_side_effects_is_kept():
|
||||
text = same_as_unoptimized("f(x) -> print(x); unused = f(3); 1")
|
||||
assert "unused=f(3)" in text
|
||||
|
||||
|
||||
def test_lambdas_are_pure_or_impure_by_their_body():
|
||||
text = same_as_unoptimized("g = (x) -> x * 2; y = g(3); 1")
|
||||
assert "y=" not in text and "g=" not in text
|
||||
text = same_as_unoptimized("g = (x) -> print(x); y = g(3); 1")
|
||||
assert "y=g(3)" in text
|
||||
# a lambda reaching another one captured when it was created
|
||||
text = same_as_unoptimized("p = (x) -> print(x); g = (x) -> p(x) + 1; h = g; y = h(2); 1")
|
||||
assert "y=h(2)" in text
|
||||
|
||||
|
||||
def test_recursive_lambdas_through_parameters():
|
||||
pure = "f = (g, n) -> (n <= 1 ? 1 : n * g(g, n - 1)); r = f(f, 4); 1"
|
||||
assert "r=" not in same_as_unoptimized(pure)
|
||||
impure = "f = (g, n) -> (n <= 1 ? print(1) : n * g(g, n - 1)); r = f(f, 4); 1"
|
||||
assert "r=f(f,4)" in same_as_unoptimized(impure)
|
||||
# a and b only reach each other through their parameters; the side effect is in b
|
||||
mutual = "a = (p, q, n) -> (n <= 0 ? 0 : q(q, p, n - 1)); b = (p, q, n) -> (n <= 0 ? print(n) : q(q, p, n - 1)); r = a(a, b, 3); 1"
|
||||
assert "r=a(a,b,3)" in same_as_unoptimized(mutual)
|
||||
assert "r=" not in same_as_unoptimized(mutual.replace("print(n)", "n"))
|
||||
|
||||
|
||||
def test_lambdas_passed_to_functions_and_unknown_lambdas():
|
||||
assert "r=h(k)" in same_as_unoptimized("h(p) -> p(3); k = (x) -> print(x); r = h(k); 1")
|
||||
assert "r=" not in same_as_unoptimized("h(p) -> p(3); k = (x) -> x * 2; r = h(k); 1")
|
||||
# a lambda returned from a function is not followed, so calls through it stay
|
||||
assert "y=f(2)" in same_as_unoptimized("mk() -> (x) -> x * 2; f = mk(); y = f(2); 1")
|
||||
|
||||
|
||||
def test_loops_keep_loop_carried_values_and_drop_dead_ones():
|
||||
text = same_as_unoptimized("s = 0; t = 0; for (i in [1, 2, 3]) { t = t + i; s = s + i; } s")
|
||||
assert "t=t+i" not in text and "t=0" not in text and "s=s+i" in text
|
||||
text = same_as_unoptimized("i = 0; j = 0; while (i < 3) { j = i * 2; i = i + 1; } i")
|
||||
assert "j=i*2" not in text and "i=i+1" in text
|
||||
text = same_as_unoptimized("acc = 0; for (v in [1, 2]) { tmp = v * 2; acc = acc + tmp; } acc")
|
||||
assert "tmp=v*2" in text
|
||||
|
||||
|
||||
def test_break_and_conditional_assignments():
|
||||
script = "x = 1; for (i in [1, 2, 3]) { if (i == 2) { x = i; break; } } x"
|
||||
text = same_as_unoptimized(script)
|
||||
assert "x=i" in text and "x=1" in text
|
||||
text = same_as_unoptimized("x = 5; if (V0 > 0) x = 2; x", {"V0": 1.0})
|
||||
assert "x=5" in text
|
||||
|
||||
|
||||
def test_compound_and_indexed_assignments_read_the_target():
|
||||
text = same_as_unoptimized("a = [1, 2, 3]; a[0] = 5; a += 1; a")
|
||||
assert "a[0]=5" in text and "a=[1,2,3]" in text
|
||||
text = same_as_unoptimized("a = [1, 2, 3]; a[0] = 5; 1")
|
||||
assert "a[0]=5" not in text and "a=[1,2,3]" not in text
|
||||
|
||||
|
||||
def test_last_statement_of_a_block_stays():
|
||||
text = same_as_unoptimized("if (1) { a = 2; b = 3; } 7")
|
||||
assert "a=2" not in text and "b=3" in text
|
||||
|
||||
|
||||
def test_unused_inputs_are_no_longer_requested():
|
||||
tree = remove_dead_code(parse_expr("u = V1 * 2; V0"))
|
||||
assert checkLazyNew(tree, {"V0": None, "V1": None}, {}) == {"V0"}
|
||||
Reference in New Issue
Block a user