Add action to set custom cost

[saulshanabrook]

Alternative to #353 that uses a new action instead of specifying a new table.

A few differences over the other approach:

• Allows overriding cost for any function, instead of only those who specify a custom cost function in the constructor. This supports the use case where one person defines a number of functions, and then another comes along and imports those and changes the costs, without having to fork the original definitions.
• Doesn't require a new function, so more succinct and obvious to use in Python (see details for API comparison).
• IMHO slightly more natural feeling API, but this is rather subjective :)

This PR API:

(datatype E
  (Foo String))

(union (Foo "x") (Foo "y"))
(union (Foo "y") (Foo "z"))

(cost (Foo "x") 17)
(cost (Foo "y") 11)
(cost (Foo "z") 15)

(extract (Foo "y"))

Other PR API:

(datatype E
  (Foo String :cost FooCost))

(union (Foo "x") (Foo "y"))
(union (Foo "y") (Foo "z"))

(set (FooCost "x") 17)
(set (FooCost "y") 11)
(set (FooCost "z") 15)

(extract (Foo "y"))

Python API comparisons

This PR:

class E(Expr):
    pass

@function
def Foo(x: StringLike) -> E: ...

egraph = EGraph()

egraph.register(
    union(Foo("x")).with_(Foo("y")),
    union(Foo("y")).with_(Foo("z")),
    cost(Foo("x")).to(i64(17)),
    cost(Foo("y")).to(i64(11)),
    cost(Foo("z")).to(i64(15)),
)

Other PR:

class E(Expr):
    pass

def FooCost(x: StringLike) -> i64: ...

@function(cost=FooCost)
def Foo(x: StringLike) -> E: ...

egraph = EGraph()

egraph.register(
    union(Foo("x")).with_(Foo("y")),
    union(Foo("y")).with_(Foo("z")),
    set_(FooCost("x")).to(i64(17)),
    set_(FooCost("y")).to(i64(11)),
    set_(FooCost("z")).to(i64(15)),
)

Closes #294

[mwillsey]

I like this approach in general. Where does the min occur? It'd be nice to see an egglog example as well as python. In my version, you'd write the cost function with :merge (min old new), so you didn't have to worry about "manually" min-ing things. Here it seems like you'd have to account for that. You could bake the min into update_cost, but then it's unclear how you "reset" the cost if that's ever something you want to do. Thoughts?

[saulshanabrook]

Thanks for taking a look Max!

I updated the description to include the egglog comparison of APIs as well.

Regarding min, I believe currently it's similar to a on merge function that just takes the new value. It does it here.

We could add a way to customize this, or have a different default. I could imagine adding the ability to do something like :cost-merge (max old new) to functions as well. We could also add a reset-cost action I suppose if we find that to be needed.

For both of these, I wonder about waiting to add them until/if they are needed for a particular use case? Or maybe there is a way to add them that is easy enough, just hesitant to add extra complexity until it's needed.

[mwillsey]

Hmm now I'm questioning whether either this or my PR provide the functionality needed from the zulip thread. We need to write out a full example that actually calculates the cost. So in your world, it would look something like this right?

(datatype Expr
  (Mul Expr Expr)
  ...)

(rule ((Mul a b))
      ((set-cost (Mul a b) (+ 10 (+ (cost?? a) (cost?? b)))))

I think we need a way to read the cost of any expression of that sort, otherwise I'm not how the rules to set cost should work.

[oflatt]

Here's how you might do it by defining a table to store the minimum cost for an eclass

(function Cost (Expr) i64 :merge (min old new))
(rule ((Mul a b)
         (MulCost a b c))
       ((set (Cost (Mul a b)) c)))

(rule ((Mul a b)
         (Cost a a-cost)
         (Cost b b-cost))
      ((set (MulCost a b) (+ 1 (+ a-cost b-cost))))) 

[mwillsey]

Right, okay that makes sense. It's quite verbose, and the first rule is pure boilerplate, it should never really look different than that. You'd also need to make rules for every other constructor (even if you were still using AST size as their cost). Maybe I'm missing something?

[oflatt]

Yeah, it's pretty annoying I agree. We could generate these rules or bake them in somehow

[mwillsey]

Ok, I propose that we close/hold both this and #353 because it seems like a realistic use case requires something like @oflatt's code sample above, and I think neither PR actually makes a big dent in the boilerplate that you'd need. We could do something with more code generation, but given that I think we don't know what parts should be generated vs left alone, we should maybe just hold off.

We do however need someway of actually doing the extraction w.r.t. different costs. I think @yihozhang proposed something like passing in pairs of (constructor costfn) to the extract command? Or maybe something more imperative like a top-level command to set the costfn for a fn? What do you y'all think?

[oflatt]

Right I'm, not sure how much should be generated.

I'm against extending the language in a way that is limited or a hack
Alternatively, we could expose the tables as read-only and allow people to write extractors in rust.
Or we can figure out a pure egglog solution that we are happy with

[saulshanabrook]

I would support waiting on an implementation until we have the use cases more clearly defined.

One possibly way to do that would be some sort of needs/design doc that at least articulates the different needs, and from there we could then see how different options achieve them most elegently.

[saulshanabrook]

Closing this for now, since we might want to refine the design space here a bit more.

[AzizZayed]

How stable is this PR to use? I'm working on a project that needs this feature but I cannot wait for the redesigns and brainstorming (although that's good, thank you for your efforts and high standards).

[saulshanabrook]

@AzizZayed I will merge in the latest commits to this branch so that it is suitable for you to try working with.

We have started merging in some things with an "unstable" prefix, so that might be an option here, to give folks a chance to play with this before committing to a final design.

[saulshanabrook]

This PR is now up to date and the action is now named unstable-cost. If you try using it and have any feedback (positive or negative) that would be useful!

I will try seeing if @tylerhou's original example (which motivated this work) could be made workable using this PR https://egraphs.zulipchat.com/#narrow/stream/328972-general/topic/How.20can.20I.20find.20the.20tree.20associated.20with.20an.20extraction.3F

[saulshanabrook]

I added the ability to set costs to be expressions, in order to support the tree join example. This is how you would implement it with this PR:

(datatype JoinTree
          (Rel String)
          (Join JoinTree JoinTree))

(function size (JoinTree) i64 :merge (min old new))


(let ra (Rel "a"))
(let rb (Rel "b"))
(let rc (Rel "c"))
(let rd (Rel "d"))
(let re (Rel "e"))
(let rf (Rel "f"))

(let query (Join (Join (Join (Join (Join ra rb) rc) rd) re) rf))

; Size
(set (size ra) 50)
(set (size rb) 200)
(set (size rc) 10)
(set (size rd) 123)
(set (size re) 10000)
(set (size rf) 1)

;; cost of relation is its size minus 1, since the string arg will have a cost of 1 as well
(rule ((= (size (Rel a)) asize))
      ((cost (Rel a) (- asize 1))))
;; cost/size of join is product of sizes
(rule ((Join a b)
       (= (size a) asize)
       (= (size b) bsize))
      ((set (size (Join a b)) (* asize bsize))
       (cost (Join a b) (* asize bsize))))

; Associativity
(rewrite (Join a b) (Join b a))

; Commutativity
(rewrite (Join (Join a b) c) (Join a (Join b c)))

(run 10000)

(extract query)
(extract (size query))

Compared with the original example:

(datatype JoinTree (Rel String) (Join JoinTree JoinTree))

(function cost (JoinTree) i64 :merge (min old new))
(function size (JoinTree) i64 :merge (min old new))


(let ra (Rel "a"))
(let rb (Rel "b"))
(let rc (Rel "c"))
(let rd (Rel "d"))
(let re (Rel "e"))
(let rf (Rel "f"))

(let query (Join (Join (Join (Join (Join ra rb) rc) rd) re) rf))

; Size
(set (size ra) 50)
(set (size rb) 200)
(set (size rc) 10)
(set (size rd) 123)
(set (size re) 10000)
(set (size rf) 1)
(rule ((Join a b) (= (size a) asize) (= (size b) bsize))
      ((set (size (Join a b)) (* asize bsize))))

; Cost
(rule ((Rel a)) ((set (cost (Rel a)) (size (Rel a)))))
(rule ((Join a b) (= (size a) asize) 
       			  (= (cost a) acost)
       			  (= (size b) bsize)
       			  (= (cost b) bcost))
      ((set (cost (Join a b)) (+ (+ (* asize bsize) acost) bcost))))


; Associativity
(rewrite (Join a b) (Join b a))

; Commutativity
(rewrite (Join (Join a b) c) (Join a (Join b c)))

(run 10000)

(extract (cost query))
(extract (size query))

The computed cost for both is the same, 123000757624.

It doesn't allow for setting costs or defining custom merge functions for costs, but as you can see here you can use a companion table to do that. The cost will always just be set as the most recent value.

[AzizZayed]

Will I be able to query a unstable-cost table, like I can do with any other function? Or is this designed differently? Like: (print-function unstable-cost 10)

[saulshanabrook]

@AzizZayed No you are not able to query it like a regular function, it only shows up in the extraction calculation or when serializing the nodes. What do you need that for?

[AzizZayed]

@AzizZayed No you are not able to query it like a regular function, it only shows up in the extraction calculation or when serializing the nodes. What do you need that for?

@saulshanabrook Just for debugging purposes. It's not very important.
Can this PR be used on the demo website BTW?

[oflatt]

Initial thoughts

[saulshanabrook]

@AzizZayed There is some concern that this won't work with an updated backend that is planning to be integrated soon, so it seems like we might hold off on merging this feature in.

I will still try to update this PR to keep it in good status.

The web UI won't work with this feature.

[saulshanabrook]

I have updated this PR with changes from main, moved cost to its own action, and changed cost back from unstable-cost to cost.

[codspeed-hq]

CodSpeed Performance Report

Merging #355 will not alter performance

_{Comparing saulshanabrook:cost-action (040325a) with main (4cae848)}

Summary

✅ 8 untouched benchmarks

🆕 2 new benchmarks

Benchmarks breakdown

	Benchmark	main	saulshanabrook:cost-action	Change
🆕	extract-costfn	N/A	402.5 µs	N/A
🆕	jointree-custom-cost	N/A	17.2 ms	N/A

[saulshanabrook]

I have updated this PR to take the minimum instead of an arbitrary cost if two outputs are combined that have different costs.

Basically, taking a step back, then PR adds some additional ability to set custom costs per function but does not allow you to set cumulative costs including children. Why? Because doing so is incompatible with the current extractor as well as all extractors in extraction-gym.

The serialized e-graph must contain an integer per node for the cost, it cannot be a function that depends on which child you select.

The motivation here is that this does allow a larger subset of extraction than before, without going full hog. The sense is that this is useful enough to support more use cases, as seen by the examples included in this PR that were built from real world use cases people brought forward. cc @ezrosent

[saulshanabrook]

Update from team meeting: We don't want to support this feature now, because we are focused on stabilizing the core over adding new features that aren't directly necessary for our immediate work.

We may consider looking at it again towards the end of this release cycle or in the next one.

We are also unsure if this API is the best one to present to users or if it could be implemented as a non-core extension.

For now, if you want to have custom costs, please extract the e-graph yourself and do your own extraction how you see fit.

If you have particular use cases that would be solved by this PR or more generally would require some form of custom costs, place post a comment on the underlying issue #294 with details of your use case. That is helpful to guide the discussion of if and how we want to support this type of feature.