Tree: various speedups

[matthiasdiener]

• make dataclass non-frozen
• use mutate() for cases where a Map is modified multiple times
• remove asserts for cases that would fail immediately anyway

---

For the following microbenchmark:

from loopy.schedule.tree import Tree

def run_t():
    tree = Tree.from_root(-1)

    for i in range(100):
        tree = tree.add_node(i, i-1)

    for i in range(100):
        tree = tree.replace_node(i, 101+i)

    for i in range(100):
        tree = tree.move_node(101+i, -1)

from timeit import timeit

print(timeit(run_t, number=1000))

, this PR results in performance improvements:

Branch	Without -O	With -O
main	0.495	0.435
tree-speedups (78c19c8)	0.408	0.347

[inducer]

Suggested change

	@dataclass(frozen=False)
	@dataclass(frozen=__debug__)

(And adapt the comment suitably?)

[matthiasdiener]

$ ./run-mypy.sh
loopy/schedule/tree.py:55: error: "frozen" argument must be a True or False literal  [literal-required]

🤷

[inducer]

inducer/pytools#272

[inducer]

Use not_none?

[matthiasdiener]

In that case, wouldn't it be easier just to leave the assert there, instead of the not_none?

[alexfikl]

Would something like this work?

parent_of_node = self.parent(node)
if parent_of_node is None:
	return 0

return 1 + self.depth(parent_of_node)

No need to hide the None check in another function (which tricks mypy) + this only does the parent lookup once.

[matthiasdiener]

That sounds good, thanks! cdc3235

[inducer]

What's the performance impact with/without when using python -O?

[matthiasdiener]

What's the performance impact with/without when using python -O?

I added a table in the first comment that compares with -O.

[alexfikl]

The mypy failures are probably due to inducer/cgen#50 (which adds py.typed, so now it gets checked).

[inducer]

Thanks!