Merge pull request #1366 from herbie-fp/batchify-series

Batchifying `run-improve!`

Author: AYadrov

src/core/alt-table.rkt

@@ -12,40 +12,39 @@
          "programs.rkt")
 
 (provide (contract-out
-          (make-alt-table (pcontext? alt? any/c . -> . alt-table?))
+          (make-alt-table (batch? pcontext? alt? any/c . -> . alt-table?))
           (atab-active-alts (alt-table? . -> . (listof alt?)))
           (atab-all-alts (alt-table? . -> . (listof alt?)))
           (atab-not-done-alts (alt-table? . -> . (listof alt?)))
-          (atab-eval-altns (alt-table? (listof alt?) context? . -> . (values any/c any/c)))
+          (atab-eval-altns (alt-table? batch? (listof alt?) context? . -> . (values any/c any/c)))
           (atab-add-altns (alt-table? (listof alt?) any/c any/c context? . -> . alt-table?))
           (atab-set-picked (alt-table? (listof alt?) . -> . alt-table?))
           (atab-completed? (alt-table? . -> . boolean?))
-          (atab-min-errors (alt-table? . -> . (listof real?)))))
+          (atab-min-errors (alt-table? . -> . (listof real?)))
+          (alt-batch-costs (batch? representation? . -> . (batchref? . -> . real?)))))
 
 ;; Public API
 
 (struct alt-table (point-idx->alts alt->point-idxs alt->done? alt->cost pcontext all) #:prefab)
 
-(define (alt-batch-cost batch brfs repr)
+(define (alt-batch-costs batch repr)
   (define node-cost-proc (platform-node-cost-proc (*active-platform*)))
-  (define costs
-    (batch-map batch
-               (λ (get-args-costs node)
-                 (match node
-                   [(? literal?) ((node-cost-proc node repr))]
-                   [(? symbol?) ((node-cost-proc node repr))]
-                   [(? number?) 0] ; specs
-                   [(approx _ impl) (get-args-costs impl)]
-                   [(list (? (negate impl-exists?) impl) args ...) 0] ; specs
-                   [(list impl args ...)
-                    (define cost-proc (node-cost-proc node repr))
-                    (define itypes (impl-info impl 'itype))
-                    (apply cost-proc (map get-args-costs args))]))))
-  (map costs brfs))
-
-(define (make-alt-table pcontext initial-alt ctx)
-  (define cost (alt-cost initial-alt (context-repr ctx)))
-  (define errs (errors (alt-expr initial-alt) pcontext ctx))
+  (batch-map batch
+             (λ (get-args-costs node)
+               (match node
+                 [(? literal?) ((node-cost-proc node repr))]
+                 [(? symbol?) ((node-cost-proc node repr))]
+                 [(? number?) 0] ; specs
+                 [(approx _ impl) (get-args-costs impl)]
+                 [(list (? (negate impl-exists?) impl) args ...) 0] ; specs
+                 [(list impl args ...)
+                  (define cost-proc (node-cost-proc node repr))
+                  (define itypes (impl-info impl 'itype))
+                  (apply cost-proc (map get-args-costs args))]))))
+
+(define (make-alt-table batch pcontext initial-alt ctx)
+  (define cost ((alt-batch-costs batch (context-repr ctx)) (alt-expr initial-alt)))
+  (define errs (batchref-errors (alt-expr initial-alt) pcontext ctx))
   (alt-table (for/vector #:length (pcontext-length pcontext)
                          ([err (in-list errs)])
                (list (pareto-point cost err (list initial-alt))))
@@ -180,10 +179,10 @@
                [alt->done? (hash-remove* alt->done? altns)]
                [alt->cost (hash-remove* alt->cost altns)]))
 
-(define (atab-eval-altns atab altns ctx)
-  (define-values (batch brfs) (progs->batch (map alt-expr altns) #:vars (context-vars ctx)))
+(define (atab-eval-altns atab batch altns ctx)
+  (define brfs (map alt-expr altns))
   (define errss (batch-errors batch brfs (alt-table-pcontext atab) ctx))
-  (define costs (alt-batch-cost batch brfs (context-repr ctx)))
+  (define costs (map (alt-batch-costs batch (context-repr ctx)) brfs))
   (values errss costs))
 
 (define (atab-add-altns atab altns errss costs ctx)
@@ -215,22 +214,26 @@
   (match-define (alt-table point-idx->alts alt->point-idxs alt->done? alt->cost pcontext _) atab)
   (define max-error (+ 1 (expt 2 (representation-total-bits (context-repr ctx)))))
 
-  (define point-idx->alts*
-    (for/vector #:length (vector-length point-idx->alts)
-                ([pcurve (in-vector point-idx->alts)]
-                 [err (in-list errs)])
-      (cond
-        [(< err max-error) ; Only include points if they are valid
-         (define ppt (pareto-point cost err (list altn)))
-         (pareto-union (list ppt) pcurve #:combine append)]
-        [else pcurve])))
-
-  (alt-table point-idx->alts*
-             (hash-set alt->point-idxs altn #f)
-             (hash-set alt->done? altn #f)
-             (hash-set alt->cost altn cost)
-             pcontext
-             #f))
+  ;; Check  whether altn is already inserted into atab
+  (match (hash-has-key? alt->point-idxs altn)
+    [#f
+     (define point-idx->alts*
+       (for/vector #:length (vector-length point-idx->alts)
+                   ([pcurve (in-vector point-idx->alts)]
+                    [err (in-list errs)])
+         (cond
+           [(< err max-error) ; Only include points if they are valid
+            (define ppt (pareto-point cost err (list altn)))
+            (pareto-union (list ppt) pcurve #:combine append)]
+           [else pcurve])))
+
+     (alt-table point-idx->alts*
+                (hash-set alt->point-idxs altn #f)
+                (hash-set alt->done? altn #f)
+                (hash-set alt->cost altn cost)
+                pcontext
+                #f)]
+    [_ atab]))
 
 (define (atab-min-errors atab)
   (define pnt-idx->alts (alt-table-point-idx->alts atab))

src/core/batch.rkt

@@ -40,6 +40,8 @@
 
 (struct batchref (batch idx) #:transparent)
 
+;; --------------------------------- CORE BATCH FUNCTION ------------------------------------
+
 (define (batch-empty)
   (batch (make-dvector) (make-hash) (make-hasheq)))
 
@@ -107,13 +109,11 @@
   (values out brfs))
 
 (define (batch->progs b brfs)
-  (brfs-belong-to-batch? b brfs)
   (map (batch-exprs b) brfs))
 
 ;; batch-map does not iterate over nodes that are not a child of brf
 ;; A lot of parts of Herbie rely on that
 (define (batch-map batch f)
-  (define len (batch-length batch))
   (define out (make-dvector))
   (define visited (make-dvector))
   (λ (brf)
@@ -131,22 +131,24 @@
          (dvector-set! visited idx #t)
          res]))))
 
-(define (batch-apply b brfs f)
-  (define out (batch-empty))
-  (define apply-f
-    (λ (remap node)
-      (define node-with-batchrefs (expr-recurse node (lambda (ref) (batchref b ref))))
-      (define node* (f node-with-batchrefs))
-      (define brf*
-        (let loop ([node* node*])
-          (match node*
-            [(? batchref? brf) (remap (batchref-idx brf))]
-            [_ (batch-push! out (expr-recurse node* (compose batchref-idx loop)))])))
-      brf*))
-  (define brfs* (map (batch-map b apply-f) brfs))
-  (values out brfs*))
+(define (batch-ref batch reg)
+  (dvector-ref (batch-nodes batch) reg))
 
-(define (batch-apply! b f)
+(define (batch-pull brf)
+  (define (unmunge brf)
+    (expr-recurse (deref brf) unmunge))
+  (unmunge brf))
+
+(define (brfs-belong-to-batch? batch brfs)
+  (unless (andmap (compose (curry equal? batch) batchref-batch) brfs)
+    (error 'brfs-belong-to-batch? "One of batchrefs does not belong to the provided batch")))
+
+;; --------------------------------- CUSTOM BATCH FUNCTION ------------------------------------
+
+;; out - batch to where write new nodes
+;; b - batch from which to read nodes
+;; f - function to modify nodes from b
+(define (batch-apply-internal out b f)
   (batch-map b
              (λ (remap node)
                (define node-with-batchrefs (expr-recurse node (lambda (ref) (batchref b ref))))
@@ -155,9 +157,22 @@
                  (let loop ([node* node*])
                    (match node*
                      [(? batchref? brf) (remap (batchref-idx brf))]
-                     [_ (batch-push! b (expr-recurse node* (compose batchref-idx loop)))])))
+                     [_ (batch-push! out (expr-recurse node* (compose batchref-idx loop)))])))
                brf*)))
 
+;; Allocates new batch
+(define (batch-apply b brfs f)
+  (define out (batch-empty))
+  (define apply-f (batch-apply-internal out b f))
+  (define brfs* (map apply-f brfs))
+  (values out brfs*))
+
+;; Modifies batch in-place
+(define (batch-apply! b f)
+  (batch-apply-internal b b f))
+
+;; Function returns indices of children nodes within a batch for given roots,
+;;   where a child node is a child of a root + meets a condition - (condition node)
 (define (batch-reachable batch brfs #:condition [condition (const #t)])
   ; Little check
   (brfs-belong-to-batch? batch brfs)
@@ -207,21 +222,11 @@
                  (apply + 1 (map get-children-counts args)))))
   (apply + (map counts brfs)))
 
-(define (brfs-belong-to-batch? batch brfs)
-  (unless (andmap (compose (curry equal? batch) batchref-batch) brfs)
-    (error 'brfs-belong-to-batch? "One of batchrefs does not belong to the provided batch")))
-
 ;; The function removes any zombie nodes from batch with respect to the brfs
 (define (batch-copy-only batch brfs)
   (batch-apply batch brfs identity))
 
-(define (batch-ref batch reg)
-  (dvector-ref (batch-nodes batch) reg))
-
-(define (batch-pull brf)
-  (define (unmunge brf)
-    (expr-recurse (deref brf) unmunge))
-  (unmunge brf))
+;; --------------------------------- TESTS ---------------------------------------
 
 ; Tests for progs->batch and batch->progs
 (module+ test

src/core/compiler.rkt

@@ -74,10 +74,9 @@
   (define vars (context-vars ctx))
   (define num-vars (length vars))
 
-  (timeline-push! 'compiler (batch-tree-size batch brfs) (batch-length batch))
-
   ; Here we need to keep vars even though no roots refer to the vars
   (define-values (batch* brfs*) (batch-for-compiler batch brfs vars))
+  (timeline-push! 'compiler (batch-tree-size batch* brfs*) (batch-length batch*))
 
   (define instructions
     (for/vector #:length (- (batch-length batch*) num-vars)

src/core/mainloop.rkt

@@ -34,6 +34,7 @@
 (define *start-prog* (make-parameter #f))
 (define *pcontext* (make-parameter #f))
 (define *preprocessing* (make-parameter '()))
+(define *global-batch* (make-parameter #f))
 
 ;; These high-level functions give the high-level workflow of Herbie:
 ;; - Initial steps: explain, preprocessing, initialize the alt table
@@ -47,9 +48,11 @@
   (define pcontext* (preprocess-pcontext context pcontext preprocessing))
   (*pcontext* pcontext*)
   (*start-prog* initial)
+  (*global-batch* (batch-empty))
   (*preprocessing* preprocessing)
-  (define start-alt (alt initial 'start '()))
-  (^table^ (make-alt-table pcontext start-alt context))
+  (define initial-brf (batch-add! (*global-batch*) initial))
+  (define start-alt (alt initial-brf 'start '()))
+  (^table^ (make-alt-table (*global-batch*) pcontext start-alt context))
 
   (for ([iteration (in-range (*num-iterations*))]
         #:break (atab-completed? (^table^)))
@@ -63,8 +66,7 @@
 
 (define (extract!)
   (timeline-push-alts! '())
-
-  (define all-alts (atab-all-alts (^table^)))
+  (define all-alts (unbatchify-alts (*global-batch*) (atab-all-alts (^table^))))
   (define joined-alts (make-regime! all-alts (*start-prog*))) ;; HERE
   (define annotated-alts (add-derivations! joined-alts))
 
@@ -111,15 +113,19 @@
 
 (define (score-alt alt)
   (errors-score (errors (alt-expr alt) (*pcontext*) (*context*))))
+(define (batch-score-alts altns)
+  (map errors-score (batch-errors (*global-batch*) (map alt-expr altns) (*pcontext*) (*context*))))
 
 ; Pareto mode alt picking
 (define (choose-mult-alts altns)
   (define repr (context-repr (*context*)))
   (cond
     [(< (length altns) (*pareto-pick-limit*)) altns] ; take max
     [else
-     (define best (argmin score-alt altns))
-     (define altns* (sort (set-remove altns best) < #:key (curryr alt-cost repr)))
+     (define scores (batch-score-alts altns))
+     (define best (list-ref altns (index-of scores (argmin identity scores))))
+     (define alt-costs (alt-batch-costs (*global-batch*) repr))
+     (define altns* (sort (set-remove altns best) < #:key (compose alt-costs alt-expr)))
      (define simplest (car altns*))
      (define altns** (cdr altns*))
      (define div-size (round (/ (length altns**) (- (*pareto-pick-limit*) 1))))
@@ -128,16 +134,18 @@
                (list-ref altns** (- (* i div-size) 1))))]))
 
 (define (timeline-push-alts! picked-alts)
+  (define exprs (batch-exprs (*global-batch*)))
   (define fresh-alts (atab-not-done-alts (^table^)))
   (define repr (context-repr (*context*)))
-  (for ([alt (atab-active-alts (^table^))])
+  (for ([alt (atab-active-alts (^table^))]
+        [sc (in-list (batch-score-alts (atab-active-alts (^table^))))])
     (timeline-push! 'alts
-                    (~a (alt-expr alt))
+                    (~a (exprs (alt-expr alt)))
                     (cond
                       [(set-member? picked-alts alt) "next"]
                       [(set-member? fresh-alts alt) "fresh"]
                       [else "done"])
-                    (score-alt alt)
+                    sc
                     (~a (representation-name repr)))))
 
 (define (choose-alts!)
@@ -150,7 +158,7 @@
   (void))
 
 ;; Converts a patch to full alt with valid history
-(define (reconstruct! global-batch alts)
+(define (reconstruct! alts)
   ;; extracts the base expressions of a patch as a batchref
   (define (get-starting-expr altn)
     (match (alt-prevs altn)
@@ -169,7 +177,7 @@
              [(list 'evaluate) (list 'evaluate loc0)]
              [(list 'taylor name var) (list 'taylor loc0 name var)]
              [(list 'rr input proof) (list 'rr loc0 input proof)]))
-         (define expr* (batch-location-set global-batch (alt-expr orig) loc0 (alt-expr altn)))
+         (define expr* (batch-location-set (*global-batch*) (alt-expr orig) loc0 (alt-expr altn)))
          (alt expr* event* (list (loop (first prevs))))])))
 
   (^patched^ (remove-duplicates
@@ -183,8 +191,6 @@
                                (reconstruct-alt altn loc full-altn))))))
               #:key (compose batchref-idx alt-expr)))
 
-  (^patched^ (unbatchify-alts global-batch (^patched^)))
-  ; No need to unmunge ^next-alts^
   (void))
 
 ;; Finish iteration
@@ -197,7 +203,7 @@
   (define orig-fresh-alts (atab-not-done-alts (^table^)))
   (define orig-done-alts (set-subtract orig-all-alts (atab-not-done-alts (^table^))))
 
-  (define-values (errss costs) (atab-eval-altns (^table^) (^patched^) (*context*)))
+  (define-values (errss costs) (atab-eval-altns (^table^) (*global-batch*) (^patched^) (*context*)))
   (timeline-event! 'prune)
   (^table^ (atab-add-altns (^table^) (^patched^) errss costs (*context*)))
   (define final-fresh-alts (atab-not-done-alts (^table^)))
@@ -232,14 +238,10 @@
   (unless (^next-alts^)
     (choose-alts!))
 
-  (define-values (global-batch brfs) (progs->batch (map alt-expr (^next-alts^))))
-  (define (make-batchref x brf)
-    (struct-copy alt x [expr brf]))
-
-  (^next-alts^ (map make-batchref (^next-alts^) brfs))
-  (define brfs* (batch-reachable global-batch brfs #:condition node-is-impl?))
+  (define brfs (map alt-expr (^next-alts^)))
+  (define brfs* (batch-reachable (*global-batch*) brfs #:condition node-is-impl?))
 
-  (reconstruct! global-batch (generate-candidates global-batch brfs*))
+  (reconstruct! (generate-candidates (*global-batch*) brfs*))
   (finalize-iter!)
   (void))