Skip to content

Commit

Permalink
Merge pull request #448 from egraphs-good/oflatt-subst-fuel
Browse files Browse the repository at this point in the history 
Substitution with fuel
  • Loading branch information
@oflatt
oflatt authored Apr 12, 2024
2 parents 4481b07 + 16a54ea commit d5ec258
Show file tree
Hide file tree
Showing 3 changed files with 154 additions and 40 deletions.
4 changes: 3 additions & 1 deletion dag_in_context/src/schedule.egg
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,9 @@
(saturate type-helpers)
type-analysis
)
(saturate subst) ;; subst depends on type information for the argument
(saturate
(saturate subst-helpers)
subst) ;; subst depends on type information for the argument
(saturate interval-analysis)
(saturate
;; make sure to saturate context-helpers, since they ensure the other rules saturate
Expand Down
95 changes: 61 additions & 34 deletions dag_in_context/src/utility/subst.egg
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,28 +3,43 @@
;; It performs the substitution, copying over the equalities from the original eclass.
;; It only places context on the leaf nodes.

(ruleset subst-helpers)
(ruleset subst)

;; (Subst assumption to in) substitutes to for `(Arg ty)` in `in`.
(let inf-fuel 1000000)

;; (Subst fuel assumption to in) substitutes to for `(Arg ty)` in `in`.
;; It also replaces any contexts found by updating them to `assumption`.
;; `assumption` *justifies* this substitution, as the context that the result is used in.
;; In other words, it must refine the equivalence relation of `in` with `to` as the argument.
(function Subst (Assumption Expr Expr) Expr :unextractable)
;; Subst rules always saturate, even when the fuel is infinite. The fuel is there to control
;; exponential substitution for optimizations like function inlining.
;; Fuel is decreased every time the Subst propagates to a child node. Subst does not continue
;; to propagate if fuel reaches 0.
(function Subst (i64 Assumption Expr Expr) Expr :unextractable)
;; (SubstLeaf to in) substitutes `to` for `in` in a leaf node `in`
;; SubstLeaf assumes context has already been added
(function SubstLeaf (Expr Expr) Expr :unextractable)
(function SubstList (Assumption Expr ListExpr) ListExpr :unextractable)
(function SubstList (i64 Assumption Expr ListExpr) ListExpr :unextractable)


;; if you find a subst with a higher fuel than yours, subsume this one
(rule ((= lower (Subst fuel1 assum to in))
(= higher (Subst fuel2 assum to in))
(> fuel2 fuel1))
((union lower higher)
(subsume (Subst fuel1 assum to in)))
:ruleset subst-helpers)

;; Base case- leaf nodes
;; leaf node without context
(rule ((= lhs (Subst assum to leaf))
(rule ((= lhs (Subst fuel assum to leaf))
(IsLeaf leaf))
((union lhs (InContext assum (SubstLeaf to leaf))))
:ruleset subst)
;; leaf node with context
;; replace this context- subst assumes the context is more specific
(rule ((= lhs (Subst assum to (InContext oldctx leaf)))
(rule ((= lhs (Subst fuel assum to (InContext oldctx leaf)))
(IsLeaf leaf))
((union lhs (InContext assum (SubstLeaf to leaf))))
:ruleset subst)
Expand All @@ -44,60 +59,72 @@


;; Operators
(rewrite (Subst assum to (Bop op c1 c2))
(Bop op (Subst assum to c1)
(Subst assum to c2))
(rewrite (Subst fuel assum to (Bop op c1 c2))
(Bop op (Subst (- fuel 1) assum to c1)
(Subst (- fuel 1) assum to c2))
:when ((> fuel 0))
:ruleset subst)
(rewrite (Subst assum to (Uop op c1))
(Uop op (Subst assum to c1))
(rewrite (Subst fuel assum to (Uop op c1))
(Uop op (Subst (- fuel 1) assum to c1))
:when ((> fuel 0))
:ruleset subst)
(rewrite (Subst assum to (Get c1 index))
(Get (Subst assum to c1) index)
(rewrite (Subst fuel assum to (Get c1 index))
(Get (Subst (- fuel 1) assum to c1) index)
:when ((> fuel 0))
:ruleset subst)
(rewrite (Subst assum to (Alloc id c1 state ty))
(Alloc id (Subst assum to c1) (Subst assum to state) ty)
(rewrite (Subst fuel assum to (Alloc id c1 state ty))
(Alloc id (Subst (- fuel 1) assum to c1) (Subst (- fuel 1) assum to state) ty)
:when ((> fuel 0))
:ruleset subst)
(rewrite (Subst assum to (Call name c1))
(Call name (Subst assum to c1))
(rewrite (Subst fuel assum to (Call name c1))
(Call name (Subst (- fuel 1) assum to c1))
:when ((> fuel 0))
:ruleset subst)


;; Tuple operators
(rewrite (Subst assum to (Single c1))
(Single (Subst assum to c1))
(rewrite (Subst fuel assum to (Single c1))
(Single (Subst (- fuel 1) assum to c1))
:when ((> fuel 0))
:ruleset subst)
(rewrite (Subst assum to (Concat c1 c2))
(Concat (Subst assum to c1) (Subst assum to c2))
(rewrite (Subst fuel assum to (Concat c1 c2))
(Concat (Subst (- fuel 1) assum to c1) (Subst (- fuel 1) assum to c2))
:when ((> fuel 0))
:ruleset subst)


;; Control flow
(rewrite (Subst assum to (Switch pred inputs branches))
(Switch (Subst assum to pred)
(Subst assum to inputs)
(rewrite (Subst fuel assum to (Switch pred inputs branches))
(Switch (Subst (- fuel 1) assum to pred)
(Subst (- fuel 1) assum to inputs)
branches)
:when ((> fuel 0))
:ruleset subst)
(rewrite (Subst assum to (If pred inputs c1 c2))
(If (Subst assum to pred)
(Subst assum to inputs)
(rewrite (Subst fuel assum to (If pred inputs c1 c2))
(If (Subst (- fuel 1) assum to pred)
(Subst (- fuel 1) assum to inputs)
c1
c2)
:when ((> fuel 0))
:ruleset subst)
(rewrite (Subst assum to (DoWhile in out))
(DoWhile (Subst assum to in)
(rewrite (Subst fuel assum to (DoWhile in out))
(DoWhile (Subst (- fuel 1) assum to in)
out)
:when ((> fuel 0))
:ruleset subst)

;; List operators
(rewrite (SubstList assum to (Cons c1 c2))
(Cons (Subst assum to c1)
(SubstList assum to c2))
(rewrite (SubstList fuel assum to (Cons c1 c2))
(Cons (Subst fuel assum to c1)
(SubstList fuel assum to c2))
:when ((> fuel 0))
:ruleset subst)
(rewrite (SubstList assum to (Nil))
(rewrite (SubstList fuel assum to (Nil))
(Nil)
:when ((> fuel 0))
:ruleset subst)

;; substitute into function (convenience)
(rewrite (Subst assum to (Function name inty outty body))
(Function name inty outty (Subst assum to body))
(rewrite (Subst fuel assum to (Function name inty outty body))
(Function name inty outty (Subst fuel assum to body))
:ruleset subst)
95 changes: 90 additions & 5 deletions dag_in_context/src/utility/subst.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@ fn test_subst_nested() -> crate::Result {

let build = format!(
"
(let substituted (Subst (NoContext)
(let substituted (Subst inf-fuel (NoContext)
{replace_with}
{expr}))"
);
Expand Down Expand Up @@ -76,7 +76,7 @@ fn test_subst_makes_new_context() -> crate::Result {
.with_arg_types(base(intt()), base(intt()));
let build = format!(
"
(let substituted (Subst (NoContext)
(let substituted (Subst inf-fuel (NoContext)
{replace_with}
{expr}))"
);
Expand Down Expand Up @@ -107,7 +107,7 @@ fn test_subst_arg_type_changes() -> crate::Result {
.with_arg_types(tupletype.clone(), base(intt()));
let build = format!(
"
(let substituted (Subst (NoContext)
(let substituted (Subst inf-fuel (NoContext)
{replace_with}
{expr}))"
);
Expand Down Expand Up @@ -151,7 +151,7 @@ fn test_subst_identity() -> crate::Result {

let build = format!(
"
(let substituted (Subst (NoContext)
(let substituted (Subst inf-fuel (NoContext)
{replace_with}
{expression}))"
);
Expand Down Expand Up @@ -183,7 +183,7 @@ fn test_subst_preserves_context() -> crate::Result {

let build = format!(
"
(let substituted (Subst (NoContext)
(let substituted (Subst inf-fuel (NoContext)
{replace_with}
{expression}))"
);
Expand All @@ -197,3 +197,88 @@ fn test_subst_preserves_context() -> crate::Result {
vec![],
)
}

// replace with, original expression, and expected for the two fuel tests
// to show that not enough fuel causes substitution to not propagate fully,
// but enough fuel causes substitution to replace the desired arg.
// For these tests, only the build/check are different but the setup is the same.
#[cfg(test)]
fn fuel_test_replacewith_expression_expected() -> (
std::rc::Rc<crate::schema::Expr>,
std::rc::Rc<crate::schema::Expr>,
std::rc::Rc<crate::schema::Expr>,
) {
use crate::ast::*;

let expression = function(
"main",
base(intt()),
base(intt()),
get(
dowhile(
single(arg()),
parallel!(less_than(get(arg(), 0), int(0)), int(3)),
),
0,
),
)
.func_with_arg_types();

let expected = function(
"main",
base(intt()),
base(intt()),
get(
dowhile(
single(inctx(noctx(), int(4))),
parallel!(less_than(get(arg(), 0), int(0)), int(3)),
),
0,
),
)
.func_with_arg_types();

let replace_with = int(4).with_arg_types(base(intt()), base(intt()));

(replace_with, expression, expected)
}

#[test]
fn test_subst_with_enough_fuel_goes() -> crate::Result {
use crate::ast::*;

let (replace_with, expression, expected) = fuel_test_replacewith_expression_expected();

let build =
format!("(let substituted (Subst 3 (NoContext) {replace_with} {expression})) (let expected {expected})");
let check = format!("(check (= substituted {expected}))");

crate::egglog_test(
&build,
&check,
vec![expression.func_to_program()],
intv(0),
intv(3),
vec![],
)
}

#[test]
fn test_subst_with_little_fuel_stops() -> crate::Result {
use crate::ast::*;

let (replace_with, expression, expected) = fuel_test_replacewith_expression_expected();

let build =
format!("(let substituted (Subst 2 (NoContext) {replace_with} {expression})) (let expected {expected})");
let check = format!("(fail (check (= substituted {expected})))");

crate::egglog_test(
&build,
&check,
vec![expression.func_to_program()],
intv(0),
intv(3),
vec![],
)
}

0 comments on commit d5ec258

Please sign in to comment.