Add action to set custom cost #355

saulshanabrook · 2024-03-12T13:57:45Z

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"))

(unstable-cost (Foo "x") 17)
(unstable-cost (Foo "y") 11)
(unstable-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)),
)

mwillsey · 2024-03-12T15:08:00Z

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 · 2024-03-12T16:41:47Z

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 · 2024-03-12T16:57:52Z

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 · 2024-03-12T18:03:39Z

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 · 2024-03-12T18:13:21Z

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 · 2024-03-12T18:22:54Z

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

mwillsey · 2024-03-13T15:45:49Z

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 · 2024-03-13T16:41:34Z

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 · 2024-03-14T13:47:48Z

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 · 2024-04-12T03:06:26Z

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

AzizZayed · 2024-08-07T04:16:44Z

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 · 2024-08-07T15:32:47Z

@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 · 2024-08-07T15:40:47Z

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 · 2024-08-07T18:24:50Z

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))
      ((unstable-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))
       (unstable-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 · 2024-08-07T19:29:37Z

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 · 2024-08-07T23:27:58Z

@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 · 2024-08-08T15:36:28Z

@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?

tests/fail-typecheck/custom-cost-wrong-type.egg

oflatt · 2024-08-08T15:44:29Z

tests/extract-costfn.egg

+(unstable-cost (Foo "y") 11)
+(unstable-cost (Foo "z") 15)
+
+(extract (Foo "y"))


This test doesn't check that we get a particular result- perhaps it should be a rust test that can?

oflatt

Initial thoughts

oflatt · 2024-08-08T15:45:07Z

src/function/table.rs

@@ -179,6 +181,7 @@ impl Table {
                    value: next,
                    timestamp: ts,
                    subsumed: subsumed || prev_subsumed,
+                    cost: cost.or(prev_cost),


shouldn't we take the smaller cost here?

I think just using the new cost also makes sense here... It should stabilize as long as the numbers it's pulling from stabilize. Like in the test that uses this, it uses the size function which has its own merge function of smallest, and then updates based on that.

oflatt · 2024-08-08T15:46:22Z

src/constraint.rs

@@ -534,15 +541,28 @@ impl CoreAction {
                    .chain(get_atom_application_constraints(f, &args, span, typeinfo)?)
                    .collect())
            }
-            CoreAction::Set(span, head, args, rhs) => {
+            CoreAction::Set(span, head, args, rhs, is_cost) => {


I think it would be cleaner to have a new core action that sets costs rather than re-use Set here.

oflatt · 2024-08-08T15:48:44Z

src/ast/desugar.rs

@@ -118,7 +118,10 @@ fn add_semi_naive_rule(desugar: &mut Desugar, rule: Rule) -> Option<Rule> {
    let mut var_set = HashSet::default();
    for head_slice in new_rule.head.0.iter_mut().rev() {
        match head_slice {
-            Action::Set(span, _, _, expr) => {
+            Action::Set(span, _, _, expr, is_cost) => {


Does cost depend only on the inputs of a function right now? Should we only allow setting costs for constructors, not functions?

Yes cost depends only on the inputs, similar to subsumed.

saulshanabrook · 2024-08-08T17:11:40Z

@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 · 2024-08-09T15:47:19Z

src/extract.rs

-        let mut cost = function.decl.cost.unwrap_or(1);
+        let mut cost = cost.unwrap_or(function.decl.cost.unwrap_or(1));


TODO: change serialization to use this cost as well

saulshanabrook requested a review from a team as a code owner March 12, 2024 13:57

saulshanabrook requested review from mwillsey and removed request for a team March 12, 2024 13:57

saulshanabrook mentioned this pull request Mar 12, 2024

Add custom extraction cost functions #353

Closed

Rename cost in eggcc so doesn't shadow action

919a246

saulshanabrook force-pushed the cost-action branch from c05cec8 to 919a246 Compare March 12, 2024 14:05

saulshanabrook closed this Apr 12, 2024

saulshanabrook mentioned this pull request Jul 26, 2024

Variable cost to variant #399

Open

Merge remote-tracking branch 'upstream/main' into cost-action

e8f7ead

saulshanabrook reopened this Aug 7, 2024

Make cost action unstable

35d2186

saulshanabrook force-pushed the cost-action branch from e1e0d79 to 35d2186 Compare August 7, 2024 15:43

saulshanabrook added 3 commits August 7, 2024 13:54

Allow costs to be set to expressions and add jointree example

8d3e864

Simplify example

45463ed

Fix costs in example

19d0005

oflatt reviewed Aug 8, 2024

View reviewed changes

tests/fail-typecheck/custom-cost-wrong-type.egg Outdated Show resolved Hide resolved

oflatt reviewed Aug 8, 2024

View reviewed changes

Fix unstable cost failed typecheck

b6e1c96

saulshanabrook commented Aug 9, 2024

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add action to set custom cost #355

Add action to set custom cost #355

saulshanabrook commented Mar 12, 2024 •

edited

Loading

mwillsey commented Mar 12, 2024

saulshanabrook commented Mar 12, 2024 •

edited

Loading

mwillsey commented Mar 12, 2024

oflatt commented Mar 12, 2024

mwillsey commented Mar 12, 2024

oflatt commented Mar 12, 2024

mwillsey commented Mar 13, 2024

oflatt commented Mar 13, 2024

saulshanabrook commented Mar 14, 2024

saulshanabrook commented Apr 12, 2024

AzizZayed commented Aug 7, 2024 •

edited

Loading

saulshanabrook commented Aug 7, 2024

saulshanabrook commented Aug 7, 2024 •

edited

Loading

saulshanabrook commented Aug 7, 2024

AzizZayed commented Aug 7, 2024

saulshanabrook commented Aug 7, 2024 •

edited

Loading

AzizZayed commented Aug 8, 2024

oflatt Aug 8, 2024

oflatt left a comment

oflatt Aug 8, 2024

saulshanabrook Aug 8, 2024

oflatt Aug 8, 2024

saulshanabrook Aug 8, 2024

oflatt Aug 8, 2024

saulshanabrook Aug 8, 2024

saulshanabrook commented Aug 8, 2024

saulshanabrook Aug 9, 2024

		let mut cost = function.decl.cost.unwrap_or(1);
		let mut cost = cost.unwrap_or(function.decl.cost.unwrap_or(1));

Add action to set custom cost #355

Are you sure you want to change the base?

Add action to set custom cost #355

Conversation

saulshanabrook commented Mar 12, 2024 • edited Loading

mwillsey commented Mar 12, 2024

saulshanabrook commented Mar 12, 2024 • edited Loading

mwillsey commented Mar 12, 2024

oflatt commented Mar 12, 2024

mwillsey commented Mar 12, 2024

oflatt commented Mar 12, 2024

mwillsey commented Mar 13, 2024

oflatt commented Mar 13, 2024

saulshanabrook commented Mar 14, 2024

saulshanabrook commented Apr 12, 2024

AzizZayed commented Aug 7, 2024 • edited Loading

saulshanabrook commented Aug 7, 2024

saulshanabrook commented Aug 7, 2024 • edited Loading

saulshanabrook commented Aug 7, 2024

AzizZayed commented Aug 7, 2024

saulshanabrook commented Aug 7, 2024 • edited Loading

AzizZayed commented Aug 8, 2024

Choose a reason for hiding this comment

oflatt left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

saulshanabrook commented Aug 8, 2024

Choose a reason for hiding this comment

saulshanabrook commented Mar 12, 2024 •

edited

Loading

saulshanabrook commented Mar 12, 2024 •

edited

Loading

AzizZayed commented Aug 7, 2024 •

edited

Loading

saulshanabrook commented Aug 7, 2024 •

edited

Loading

saulshanabrook commented Aug 7, 2024 •

edited

Loading