Merge pull request #987 from herbie-fp/artem-batch-egg-input

Batch-input to egg

Author: pavpanchekha

src/core/batch.rkt

@@ -13,7 +13,9 @@
          deref ; Batchref -> Expr
          batch-replace ; Batch -> (Expr<Batchref> -> Expr<Batchref>) -> Batch
          egg-nodes->batch ; Nodes -> Spec-maps -> Batch -> (Listof Batchref)
-         batchref->expr) ; Batchref -> Expr
+         batchref->expr ; Batchref -> Expr
+         batch-extract-exprs ; Batch -> (Listof Root) -> (Listof Expr)
+         remove-zombie-nodes) ; Batch -> Batch
 
 ;; This function defines the recursive structure of expressions
 (define (expr-recurse expr f)
@@ -83,6 +85,14 @@
     (timeline-push! 'compiler size (batch-length final)))
   final)
 
+(define (batch-extract-exprs b roots)
+  (define exprs (make-vector (batch-length b)))
+  (for ([node (in-vector (batch-nodes b))]
+        [idx (in-naturals)])
+    (vector-set! exprs idx (expr-recurse node (lambda (x) (vector-ref exprs x)))))
+  (for/list ([root roots])
+    (vector-ref exprs root)))
+
 (define (batch->progs b)
   (define exprs (make-vector (batch-length b)))
   (for ([node (in-vector (batch-nodes b))]
@@ -111,7 +121,9 @@
   (define roots (vector-map (curry vector-ref mapping) (batch-roots b)))
   (mutable-batch->batch out roots))
 
-; The function removes any zombie nodes from batch
+; The function removes any zombie nodes from batch with respect to the roots
+; Time complexity: O(|R| + |N|), where |R| - number of roots, |N| - length of nodes
+; Space complexity: O(|N| + |N*| + |R|), where |N*| is a length of nodes without zombie nodes
 (define (remove-zombie-nodes input-batch)
   (define nodes (batch-nodes input-batch))
   (define roots (batch-roots input-batch))
@@ -123,29 +135,21 @@
   (for ([node (in-vector nodes (- nodes-length 1) -1 -1)]
         [zmb (in-vector zombie-mask (- nodes-length 1) -1 -1)]
         #:when (not zmb))
-    (match node
-      [(list op args ...) (map (λ (n) (vector-set! zombie-mask n #f)) args)]
-      [(approx spec impl) (vector-set! zombie-mask impl #f)]
-      [_ void]))
+    (expr-recurse node (λ (n) (vector-set! zombie-mask n #f))))
 
-  (define mappings (build-vector nodes-length values))
+  (define mappings (make-vector nodes-length -1))
+  (define (remap idx)
+    (vector-ref mappings idx))
 
-  (define nodes* '())
+  (define out (make-mutable-batch))
   (for ([node (in-vector nodes)]
         [zmb (in-vector zombie-mask)]
-        [n (in-naturals)])
-    (if zmb
-        (for ([i (in-range n nodes-length)])
-          (vector-set! mappings i (sub1 (vector-ref mappings i))))
-        (set! nodes*
-              (cons (match node
-                      [(list op args ...) (cons op (map (curry vector-ref mappings) args))]
-                      [(approx spec impl) (approx spec (vector-ref mappings impl))]
-                      [_ node])
-                    nodes*))))
-  (set! nodes* (list->vector (reverse nodes*)))
+        [n (in-naturals)]
+        #:unless zmb)
+    (vector-set! mappings n (batch-push! out (expr-recurse node remap))))
+
   (define roots* (vector-map (curry vector-ref mappings) roots))
-  (batch nodes* roots* (batch-vars input-batch)))
+  (mutable-batch->batch out roots*))
 
 (define (batch-ref batch reg)
   (define (unmunge reg)
@@ -249,6 +253,7 @@
   (define (zombie-test #:nodes nodes #:roots roots)
     (define in-batch (batch nodes roots '()))
     (define out-batch (remove-zombie-nodes in-batch))
+    (check-equal? (batch->progs out-batch) (batch->progs in-batch))
     (batch-nodes out-batch))
 
   (check-equal? (vector 0 '(sqrt 0) 2 '(pow 2 1))

src/core/egg-herbie.rkt

@@ -73,8 +73,7 @@
                [egraph-pointer (egraph_copy (egraph-data-egraph-pointer eg-data))]))
 
 ; Adds expressions returning the root ids
-; TODO: take a batch rather than list of expressions
-(define (egraph-add-exprs egg-data exprs ctx)
+(define (egraph-add-exprs egg-data batch roots ctx)
   (match-define (egraph-data ptr herbie->egg-dict egg->herbie-dict id->spec) egg-data)
 
   ; lookups the egg name of a variable
@@ -125,39 +124,86 @@
       [(? symbol? x) (egraph_add_node ptr (symbol->string x) 0-vec root?)]
       [(? number? n) (egraph_add_node ptr (number->string n) 0-vec root?)]))
 
+  ; The function recurses on spec
+  (define (batch-parse-approx batch)
+    (batch-replace batch
+                   (lambda (node)
+                     (match node
+                       [(approx spec impl) (list '$approx spec impl)]
+                       [_ node]))))
+
+  (set-batch-roots! batch roots) ; make sure that we work with the right roots
+  ; the algorithm may crash if batch-length is zero
+  (define insert-batch
+    (if (zero? (batch-length batch)) batch (remove-zombie-nodes (batch-parse-approx batch))))
+
+  (define mappings (build-vector (batch-length insert-batch) values))
+  (define (remap x)
+    (vector-ref mappings x))
+
+  ; Inserting nodes bottom-up
+  (define root-mask (make-vector (batch-length insert-batch) #f))
+  (for ([root (in-vector (batch-roots insert-batch))])
+    (vector-set! root-mask root #t))
+  (for ([node (in-vector (batch-nodes insert-batch))]
+        [root? (in-vector root-mask)]
+        [n (in-naturals)])
+    (define node*
+      (match node
+        [(literal v _) v]
+        [(? number?) node]
+        [(? symbol?) (normalize-var node)]
+        [(list '$approx spec impl)
+         (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)))
+                      (cons spec* type))) ; preserved spec and type for extraction
+         (list '$approx (remap spec) (remap impl))]
+        [(list op (app remap args) ...) (cons op args)]))
+
+    (vector-set! mappings n (insert-node! node* root?)))
+
+  ;------------------------- DEBUGGING
   ; expr -> id
   ; expression cache
-  (define expr->id (make-hash))
+  #;(define expr->id (make-hash))
 
   ; expr -> natural
   ; inserts an expresison into the e-graph, returning its e-class id.
-  (define (insert! expr [root? #f])
-    ; transform the expression into a node pointing
-    ; to its child e-classes
-    (define node
-      (match expr
-        [(? number?) expr]
-        [(? symbol?) (normalize-var expr)]
-        [(literal v _) v]
-        [(approx spec impl)
-         (define spec* (insert! spec))
-         (define impl* (insert! impl))
-         (hash-ref! id->spec
-                    spec*
-                    (lambda ()
-                      (define spec* (normalize-spec spec)) ; preserved spec for extraction
-                      (define type (representation-type (repr-of impl ctx))) ; track type of spec
-                      (cons spec* type)))
-         (list '$approx spec* impl*)]
-        [(list op args ...) (cons op (map insert! args))]))
-    ; always insert the node if it is a root since
-    ; the e-graph tracks which nodes are roots
-    (cond
-      [root? (insert-node! node #t)]
-      [else (hash-ref! expr->id node (lambda () (insert-node! node #f)))]))
-
-  (for/list ([expr (in-list exprs)])
-    (insert! expr #t)))
+  #;(define (insert! expr [root? #f])
+      ; transform the expression into a node pointing
+      ; to its child e-classes
+      (define node
+        (match expr
+          [(literal v _) v]
+          [(? number?) expr]
+          [(? symbol?) (normalize-var expr)]
+          [(list '$approx spec impl)
+           (define spec* (insert! (vector-ref nodes spec)))
+           (define impl* (insert! (vector-ref nodes impl)))
+           (hash-ref! id->spec
+                      spec*
+                      (lambda ()
+                        (define spec* (normalize-spec (batch-ref insert-batch spec)))
+                        (define type (representation-type (repr-of-node insert-batch impl ctx)))
+                        (cons spec* type)))
+           (list '$approx spec* impl*)]
+          [(list op args ...) (cons op (map insert! (map (curry vector-ref nodes) args)))]))
+      ; always insert the node if it is a root since
+      ; the e-graph tracks which nodes are roots
+      (cond
+        [root? (insert-node! node #t)]
+        [else (hash-ref! expr->id node (lambda () (insert-node! node #f)))]))
+
+  #;(define nodes (batch-nodes insert-batch))
+  #;(for/list ([root (in-vector (batch-roots insert-batch))])
+      (insert! (vector-ref nodes root) #t))
+  ; ---------------------- END OF DEBUGGING
+
+  (for/list ([root (in-vector (batch-roots insert-batch))])
+    (remap root)))
 
 ;; runs rules on an egraph (optional iteration limit)
 (define (egraph-run egraph-data ffi-rules node-limit iter-limit scheduler const-folding?)
@@ -226,7 +272,8 @@
   (egraph_find (egraph-data-egraph-pointer egraph-data) id))
 
 (define (egraph-expr-equal? egraph-data expr goal ctx)
-  (match-define (list id1 id2) (egraph-add-exprs egraph-data (list 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))
   (= id1 id2))
 
 ;; returns a flattened list of terms or #f if it failed to expand the proof due to budget
@@ -1198,12 +1245,12 @@
            (loop (sub1 num-iters)))]
       [else (values egg-graph iteration-data)])))
 
-(define (egraph-run-schedule exprs schedule ctx)
+(define (egraph-run-schedule batch roots schedule ctx)
   ; allocate the e-graph
   (define egg-graph (make-egraph))
 
   ; insert expressions into the e-graph
-  (define root-ids (egraph-add-exprs egg-graph exprs ctx))
+  (define root-ids (egraph-add-exprs egg-graph batch roots ctx))
 
   ; run the schedule
   (define egg-graph*
@@ -1235,7 +1282,7 @@
 
 ;; Herbie's version of an egg runner.
 ;; Defines parameters for running rewrite rules with egg
-(struct egg-runner (exprs reprs schedule ctx)
+(struct egg-runner (batch roots reprs schedule ctx)
   #:transparent ; for equality
   #:methods gen:custom-write ; for abbreviated printing
   [(define (write-proc alt port mode)
@@ -1252,7 +1299,7 @@
 ;;     - scheduler: `(scheduler . <name>)` [default: backoff]
 ;;        - `simple`: run all rules without banning
 ;;        - `backoff`: ban rules if the fire too much
-(define (make-egg-runner exprs reprs schedule #:context [ctx (*context*)])
+(define (make-egg-runner batch roots reprs schedule #:context [ctx (*context*)])
   (define (oops! fmt . args)
     (apply error 'verify-schedule! fmt args))
   ; verify the schedule
@@ -1273,7 +1320,7 @@
            [_ (oops! "in instruction `~a`, unknown parameter `~a`" instr param)]))]
       [_ (oops! "expected `(<rules> . <params>)`, got `~a`" instr)]))
   ; make the runner
-  (egg-runner exprs reprs schedule ctx))
+  (egg-runner batch roots reprs schedule ctx))
 
 ;; Runs egg using an egg runner.
 ;;
@@ -1285,7 +1332,10 @@
   ;; Run egg using runner
   (define ctx (egg-runner-ctx runner))
   (define-values (root-ids egg-graph)
-    (egraph-run-schedule (egg-runner-exprs runner) (egg-runner-schedule runner) ctx))
+    (egraph-run-schedule (egg-runner-batch runner)
+                         (egg-runner-roots runner)
+                         (egg-runner-schedule runner)
+                         ctx))
   ; Perform extraction
   (match cmd
     [`(single . ,extractor) ; single expression extraction

src/core/localize.rkt

@@ -38,8 +38,10 @@
   (define lowering-rules (platform-lowering-rules))
 
   ; egg runner (2-phases for real rewrites and implementation selection)
+  (define batch (progs->batch progs))
   (define runner
-    (make-egg-runner progs
+    (make-egg-runner batch
+                     (batch-roots batch)
                      reprs
                      `((,lifting-rules . ((iteration . 1) (scheduler . simple)))
                        (,rules . ((node . ,(*node-limit*))))

src/core/mainloop.rkt

@@ -19,7 +19,8 @@
          "preprocess.rkt"
          "programs.rkt"
          "../utils/timeline.rkt"
-         "soundiness.rkt")
+         "soundiness.rkt"
+         "batch.rkt")
 (provide run-improve!)
 
 ;; The Herbie main loop goes through a simple iterative process:
@@ -374,7 +375,8 @@
      ; egg runner
      (define exprs (map alt-expr alts))
      (define reprs (map (lambda (expr) (repr-of expr (*context*))) exprs))
-     (define runner (make-egg-runner exprs reprs schedule))
+     (define batch (progs->batch exprs))
+     (define runner (make-egg-runner batch (batch-roots batch) reprs schedule))
 
      ; run egg
      (define simplified

src/core/patch.rkt

@@ -40,7 +40,8 @@
         `((,lowering-rules . ((iteration . 1) (scheduler . simple))))))
 
   ; run egg
-  (define runner (make-egg-runner (map alt-expr approxs) reprs schedule))
+  (define batch (progs->batch (map alt-expr approxs)))
+  (define runner (make-egg-runner batch (batch-roots batch) reprs schedule))
   (define simplification-options
     (simplify-batch runner
                     (typed-egg-extractor
@@ -89,7 +90,7 @@
                 [altn (in-list altns)]
                 [fv (in-list free-vars)]
                 #:when (member var fv)) ; check whether var exists in expr at all
-            (for ([_ (in-range (*taylor-order-limit*))])
+            (for ([i (in-range (*taylor-order-limit*))])
               (define gen (genexpr))
               (unless (spec-has-nan? gen)
                 (sow (alt gen `(taylor ,name ,var) (list altn) '())))))
@@ -133,7 +134,8 @@
   (define exprs (map alt-expr altns))
   (define reprs (map (curryr repr-of (*context*)) exprs))
   (timeline-push! 'inputs (map ~a exprs))
-  (define runner (make-egg-runner exprs reprs schedule #:context (*context*)))
+  (define batch (progs->batch exprs))
+  (define runner (make-egg-runner batch (batch-roots batch) reprs schedule #:context (*context*)))
   ; batchrefss is a (listof (listof batchref))
   (define batchrefss (run-egg runner `(multi . ,extractor)))
 

src/core/preprocess.rkt

@@ -13,7 +13,8 @@
          "programs.rkt"
          "points.rkt"
          "../utils/timeline.rkt"
-         "../utils/float.rkt")
+         "../utils/float.rkt"
+         "batch.rkt")
 
 (provide find-preprocessing
          preprocess-pcontext
@@ -66,7 +67,8 @@
       (,lowering-rules . ((iteration . 1) (scheduler . simple)))))
 
   ; egg query
-  (define runner (make-egg-runner (list expr) (list (context-repr ctx)) schedule))
+  (define batch (progs->batch (list expr)))
+  (define runner (make-egg-runner batch (batch-roots batch) (list (context-repr ctx)) schedule))
 
   ; run egg
   (define simplified
@@ -100,8 +102,11 @@
 
   ;; make egg runner
   (define rules (real-rules (*simplify-rules*)))
+
+  (define batch (progs->batch specs))
   (define runner
-    (make-egg-runner specs
+    (make-egg-runner batch
+                     (batch-roots batch)
                      (map (lambda (_) (context-repr ctx)) specs)
                      `((,rules . ((node . ,(*node-limit*)))))))