Merge pull request #1299 from herbie-fp/codex/use-roots-from-batch-in-make-egraph-and-make-egglog-runner

Refactor roots handling

Author: pavpanchekha

src/core/egg-herbie.rkt

@@ -67,7 +67,7 @@
                [egraph-pointer (egraph_copy (egraph-data-egraph-pointer eg-data))]))
 
 ; Adds expressions returning the root ids
-(define (egraph-add-exprs egg-data batch roots ctx)
+(define (egraph-add-exprs egg-data batch ctx)
   (match-define (egraph-data ptr id->spec) egg-data)
 
   ; normalizes an approx spec
@@ -108,13 +108,12 @@
       [(list op ids ...) (egraph_add_node ptr (~s op) (list->u32vec ids))]
       [(? (disjoin symbol? number?) x) (egraph_add_node ptr (~s x) 0-vec)]))
 
-  (define insert-batch (batch-remove-zombie batch roots))
-  (define mappings (build-vector (batch-length insert-batch) values))
+  (define mappings (build-vector (batch-length batch) values))
   (define (remap x)
     (vector-ref mappings x))
 
   ; Inserting nodes bottom-up
-  (for ([node (in-vector (batch-nodes insert-batch))]
+  (for ([node (in-vector (batch-nodes batch))]
         [n (in-naturals)])
     (define idx
       (match node
@@ -125,20 +124,20 @@
         [(approx spec impl) (insert-node! (list '$approx (remap spec) (remap impl)))]
         [(list op (app remap args) ...) (insert-node! (cons op args))]))
     (vector-set! mappings n idx))
-  (for ([root (in-vector (batch-roots insert-batch))])
+  (for ([root (in-vector (batch-roots batch))])
     (egraph_add_root ptr (remap root)))
 
-  (for ([node (in-vector (batch-nodes insert-batch))]
+  (for ([node (in-vector (batch-nodes batch))]
         #:when (approx? node))
     (match-define (approx spec impl) node)
     (hash-ref! id->spec
                (remap spec)
                (lambda ()
-                 (define spec* (normalize-spec (batch-ref insert-batch spec)))
-                 (define type (representation-type (repr-of-node insert-batch impl ctx)))
+                 (define spec* (normalize-spec (batch-ref batch spec)))
+                 (define type (representation-type (repr-of-node batch impl ctx)))
                  (cons spec* type))))
 
-  (for/list ([root (in-vector (batch-roots insert-batch))])
+  (for/list ([root (in-vector (batch-roots batch))])
     (remap root)))
 
 ;; runs rules on an egraph (optional iteration limit)
@@ -202,7 +201,7 @@
 
 (define (egraph-expr-equal? egraph-data expr goal ctx)
   (define batch (progs->batch (list expr goal)))
-  (match-define (list id1 id2) (egraph-add-exprs egraph-data batch (batch-roots batch) ctx))
+  (match-define (list id1 id2) (egraph-add-exprs egraph-data batch ctx))
   (= id1 id2))
 
 ;; returns a flattened list of terms or #f if it failed to expand the proof due to budget
@@ -1222,12 +1221,12 @@
            (loop (sub1 num-iters)))]
       [else (values egg-graph iteration-data)])))
 
-(define (egraph-run-schedule batch roots schedule ctx)
+(define (egraph-run-schedule batch schedule ctx)
   ; allocate the e-graph
   (define egg-graph (make-egraph-data))
 
   ; insert expressions into the e-graph
-  (define root-ids (egraph-add-exprs egg-graph batch roots ctx))
+  (define root-ids (egraph-add-exprs egg-graph batch ctx))
 
   ; run the schedule
   (define egg-graph*
@@ -1266,7 +1265,7 @@
 
 ;; Herbie's version of an egg runner.
 ;; Defines parameters for running rewrite rules with egg
-(struct egg-runner (batch roots reprs schedule ctx new-roots egg-graph)
+(struct egg-runner (batch reprs schedule ctx new-roots egg-graph)
   #:transparent ; for equality
   #:methods gen:custom-write ; for abbreviated printing
   [(define (write-proc alt port mode)
@@ -1282,7 +1281,7 @@
 ;;     - scheduler: `(scheduler . <name>)` [default: backoff]
 ;;        - `simple`: run all rules without banning
 ;;        - `backoff`: ban rules if the fire too much
-(define (make-egraph batch roots reprs schedule ctx)
+(define (make-egraph batch reprs schedule ctx)
   (define (oops! fmt . args)
     (apply error 'verify-schedule! fmt args))
   ; verify the schedule
@@ -1306,10 +1305,10 @@
            [_ (oops! "in instruction `~a`, unknown parameter `~a`" instr param)]))]
       [_ (oops! "expected `(<rules> . <params>)`, got `~a`" instr)]))
 
-  (define-values (root-ids egg-graph) (egraph-run-schedule batch roots schedule ctx))
+  (define-values (root-ids egg-graph) (egraph-run-schedule batch schedule ctx))
 
   ; make the runner
-  (egg-runner batch roots reprs schedule ctx root-ids egg-graph))
+  (egg-runner batch reprs schedule ctx root-ids egg-graph))
 
 (define (regraph-dump regraph root-ids reprs)
   (define dump-dir "dump-egg")

src/core/egglog-herbie-tests.rkt

@@ -363,8 +363,6 @@
       #s(approx (- (sin (+ x 1)) (sin x)) #s(hole binary64 (- (sin (- 1 (* -1 x))) (sin x))))
       #s(approx (sin (+ x 1)) #s(hole binary64 (sin (- 1 (* -1 x))))))))
 
-  (define roots (batch-roots batch))
-
   (define ctx (make-debug-context '(x eps)))
 
   (define reprs (make-list (vector-length (batch-roots batch)) (context-repr ctx)))
@@ -376,4 +374,4 @@
       (lower . ((iteration . 1) (scheduler . simple)))))
 
   (when (find-executable-path "egglog")
-    (run-egglog-multi-extractor (egglog-runner batch roots reprs schedule ctx) batch)))
+    (run-egglog-multi-extractor (egglog-runner batch reprs schedule ctx) batch)))

src/core/egglog-herbie.rkt

@@ -48,7 +48,7 @@
 
 ;; Herbie's version of an egglog runner.
 ;; Defines parameters for running rewrite rules with egglog
-(struct egglog-runner (batch roots reprs schedule ctx)
+(struct egglog-runner (batch reprs schedule ctx)
   #:transparent ; for equality
   #:methods gen:custom-write ; for abbreviated printing
   [(define (write-proc alt port mode)
@@ -65,7 +65,7 @@
 ;;     - scheduler: `(scheduler . <name>)` [default: backoff]
 ;;        - `simple`: run all rules without banning
 ;;        - `backoff`: ban rules if the fire too much
-(define (make-egglog-runner batch roots reprs schedule ctx)
+(define (make-egglog-runner batch reprs schedule ctx)
   (define (oops! fmt . args)
     (apply error 'verify-schedule! fmt args))
   ; verify the schedule
@@ -91,12 +91,11 @@
       [_ (oops! "expected `(<rules> . <params>)`, got `~a`" instr)]))
 
   ; make the runner
-  (egglog-runner batch roots reprs schedule ctx))
+  (egglog-runner batch reprs schedule ctx))
 
 ;; Runs egglog using an egglog runner by extracting multiple variants
 (define (run-egglog-multi-extractor runner output-batch) ; multi expression extraction
-  (define insert-batch
-    (batch-remove-zombie (egglog-runner-batch runner) (egglog-runner-roots runner)))
+  (define insert-batch (egglog-runner-batch runner))
   (define curr-program (make-egglog-program))
 
   ;; Dump-file

src/core/patch.rkt

@@ -76,16 +76,17 @@
   (define schedule `((lower . ((iteration . 1) (scheduler . simple)))))
 
   ; run egg
-  (define exprs (map (compose debatchref alt-expr) altns))
-  (define input-batch (progs->batch exprs))
-
   (define roots (list->vector (map (compose batchref-idx alt-expr) altns)))
-  (define reprs (map (curryr repr-of (*context*)) exprs))
+  (define reprs
+    (for/list ([root (in-vector roots)])
+      (repr-of-node global-batch root (*context*))))
+
+  (define batch* (batch-remove-zombie global-batch roots))
 
   (define runner
     (if (flag-set? 'generate 'egglog)
-        (make-egglog-runner input-batch (batch-roots input-batch) reprs schedule (*context*))
-        (make-egraph global-batch roots reprs schedule (*context*))))
+        (make-egglog-runner batch* reprs schedule (*context*))
+        (make-egraph batch* reprs schedule (*context*))))
 
   (define batchrefss
     (if (flag-set? 'generate 'egglog)
@@ -154,18 +155,17 @@
           `(,rules . ((node . ,(*node-limit*))))
           `(lower . ((iteration . 1) (scheduler . simple)))))
 
-  ; run egg
-  (define exprs (map (compose debatchref alt-expr) altns))
-  (define input-batch (progs->batch exprs))
-
   (define roots (list->vector (map (compose batchref-idx alt-expr) altns)))
-  (define reprs (map (curryr repr-of (*context*)) exprs))
-  (timeline-push! 'inputs (map ~a exprs))
+  (define reprs
+    (for/list ([root (in-vector roots)])
+      (repr-of-node global-batch root (*context*))))
+  (define batch* (batch-remove-zombie global-batch roots))
+  (timeline-push! 'inputs (map (compose ~a debatchref alt-expr) altns))
 
   (define runner
     (if (flag-set? 'generate 'egglog)
-        (make-egglog-runner input-batch (batch-roots input-batch) reprs schedule (*context*))
-        (make-egraph global-batch roots reprs schedule (*context*))))
+        (make-egglog-runner batch* reprs schedule (*context*))
+        (make-egraph batch* reprs schedule (*context*))))
 
   (define batchrefss
     (if (flag-set? 'generate 'egglog)

src/core/preprocess.rkt

@@ -72,7 +72,6 @@
   (define batch (progs->batch (cons spec (map cdr identities))))
   (define runner
     (make-egraph batch
-                 (batch-roots batch)
                  (make-list (vector-length (batch-roots batch)) (context-repr ctx))
                  `((,rules . ((node . ,(*node-limit*)))))
                  ctx))