Let-inlining optimization

tree_unique_args/src/let_inline.egg

@@ -0,0 +1,66 @@
+(ruleset let-inline)
+
+;; To implement
+;;   arg-used-Expr-count
+;;   arg-used-Expr <== easy
+;;   MapArgOverVec
+;;   set->vec
+;;   BuildSubstMapFrom
+;;   SubsetOperandMap
+
+(relation ExprUsesArgs-demand (Expr))
+(rule (
+    (= expr (Let id (All inp-ord inp-exprs) out))
+    (= inp (ListExpr-ith inp-exprs ith))
+    (ExprIsValid inp)
+    (ExprIsPure inp)
+) (
+    (arg-used-Expr-demand inp)
+))
+
+(function ExprUsesArgs (Expr) I64Set :merge (set-union old new))
+(rule (
+    (BodyContainsExpr body (Get (Arg id) i))
+) (
+    (set (arg-used-Expr e) (set-of i))
+))
+
+(rule (
+    (= expr (Let id (All inp-ord inp-exprs) out))
+    (= inp (ListExpr-ith inp-exprs ith))
+    (>= (Expr-size inp) 3) ;; so that `inp` involves at least some binary operation
+    (ExprIsPure inp)
+
+    ;; NB: arg-used-count might undercount 
+    ;; # argument is used
+    (= (arg-used-Expr-count out ith) 1)
+
+    (= inp-args (arg-used-Expr inp))
+) (
+    (let new-id (fresh-id!))
+
+    ;; Step 1: build the new inps: remove (Arg i) and 
+    ;; introduce arguments used by inp-args
+    (let new-inp-args-i64 (set->vec inp-args))
+    (let new-inp-args (MapArgOverVec new-inp-args-i64 (Arg new-id)))
+    (let new-inp-exprs (Concat (Remove in-exprs i) new-in-args))
+
+    ;; Step 2: build the pulled-in inp by substituting contexts
+    (let subst-map (BuildSubstMapFrom new-inp-args-i64 (vec-length inps))) ;; need to right shift by vec-length
+    (let new-inp (SubstOperandMap inp subst-map))
+
+    ;; Step 3: build the new outputs
+    ;; (All (Parallel) '(
+    ;;  (Get (Arg new-id) 0)
+    ;;  ...
+    ;;  (Get (Arg new-id) k-1)
+    ;;  substed-ink
+    ;;  (Get (Arg new-id) k) ;; this actually has index k+1
+    ;;  ...
+    ;;  (Get (Arg new-id) n)))
+    (let new-args (All (Parallel) (BuildLetInlineSubstArgs id new-inp k)))
+    (let new-out (DeepCopy (SubstExpr out new-args) new-id))
+
+    (let new-expr (Let new-id new-inp-exprs new-out))
+    (union expr new-expr)
+) :ruleset let-inline)

tree_unique_args/src/util.egg

@@ -2,7 +2,8 @@
 (function ListExpr-ith (ListExpr i64) Expr :unextractable)
 (function ListExpr-suffix (ListExpr i64) ListExpr :unextractable)
 (function Append (ListExpr Expr) ListExpr :unextractable)
-
+(function Concat (ListExpr ListExpr) ListExpr :unextractable)
+(function Remove (ListExpr i64) ListExpr :unextractable)
 (rule ((All order top)) ((union (ListExpr-suffix top 0) top)) :ruleset always-run)
 
 (rule ((= (ListExpr-suffix top n) (Cons hd tl)))
@@ -19,5 +20,19 @@
     (Cons e (Nil))
     :ruleset always-run)
 
+(rewrite (Concat (Cons a b) e)
+    (Cons a (Concat b e))
+    :ruleset always-run)
+(rewrite (Concat (Nil) e) e
+    :ruleset always-run)
+
+(rewrite (Remove (Cons a b) 0)
+    b
+    :ruleset always-run)
+(rewrite (Remove (Cons a b) i)
+    (Remove b (- i 1))
+    :when ((> i 0))
+    :ruleset always-run)
+
 (function Expr-size (Expr) i64 :merge (min old new))
 (function ListExpr-size (ListExpr) i64 :merge (min old new))