egraphs-good / egglog

egraphs + datalog!
https://egraphs-good.github.io/egglog/
MIT License
422 stars 46 forks source link

Built In Collection Constructor Type Checking Prohibitively Slow #388

Closed saulshanabrook closed 2 months ago

saulshanabrook commented 3 months ago

Hey all,

I tried switching some of the Python numba code to use Vec and vec-of to construct lists and found the performance degraded by a ton. After doing some profiling, I saw all the time was in type checking, in particular typechecking XOR constraints. Which makes sense because I have a few different types of vecs and it has to infer what type it is based on the contents.

I also noticed that most of the time was spent cloning the assignments hashmap, which happens when handling the XOR:

https://github.com/egraphs-good/egglog/blob/92b3d9d04fb8ecde2194578db7aa480fd6217429/src/constraint.rs#L128-L131

I thought, hm, if cloning is so expensive, then maybe we should try a copy on write hashmap, so that cloning is O(1)? I gave this a go and saw that the benchmark I was running (this is total time of Python + Rust) went from 1143.05 seconds to 10.92 seconds! A 100x speedup?

But then I noticed that on code that doesn't use a lot of vec collections, switching to the immutable data structure slows things down by a few %.

So for now, I'll probably just avoid vec-of, but I wanted to open this issue just to warn folks if they are using it when they have multiple types of vec's registered, that it maybe be really slow.

There is likely a way to solve this with some refactoring of the type checking logic.

oflatt commented 3 months ago

Great find! We should try to address this issue, imo the small % slowdown with switching to an immutable data structure is worth it if it fixes the problem.

In the long term, my opinion is still that the constraint-based typechecker is way too complex, and we should design a more principled type system that can be checked efficiently.

saulshanabrook commented 3 months ago

I had a hunch that my issue was due to my large expression sizes. I would add a big (tens of thousands of expressions) expression tree to egglog, all as one expression. This is relatively inefficient, since there are a lot of duplicated nodes.

I just added CSE to the Python egglog bindings transparently (https://github.com/egraphs-good/egglog-python/pull/179) creating let bindings automatically when a user inserts an expression to the e-graph, so that the type checking is less stressed out.

This was a huge performance improvement regardless and now with this change, I see no improvement with immutable data structures when I am using the builtin collections. So I'll close this issue and can continue using vec-of and the like without a performance issue!

yihozhang commented 3 months ago

I think it would still be useful to fix the original issue though- since there might be users who don't use the python binding and generate big S-expressions. Maybe we should reopen this issue?

saulshanabrook commented 3 months ago

Sure