Merge pull request #1298 from herbie-fp/codex/define-pcontext-in-mainloop-and-pass-as-argument

Refactor pcontext handling

Author: pavpanchekha

src/api/sandbox.rkt

@@ -107,14 +107,13 @@
   (unless pcontext
     (error 'explain "cannot run without a pcontext"))
 
-  (*pcontext* pcontext)
   (define-values (fperrors
                   sorted-explanations-table
                   confusion-matrix
                   maybe-confusion-matrix
                   total-confusion-matrix
                   freqs)
-    (explain (test-input test) (*context*) (*pcontext*)))
+    (explain (test-input test) (*context*) pcontext))
 
   sorted-explanations-table)
 
@@ -127,8 +126,7 @@
   (unless pcontext
     (error 'get-local-error "cannnot run without a pcontext"))
 
-  (*pcontext* pcontext)
-  (local-error-as-tree (test-input test) (*context*)))
+  (local-error-as-tree (test-input test) (*context*) pcontext))
 
 (define (get-sample test)
   (random) ;; Tick the random number generator, for backwards compatibility

src/core/bsearch.rkt

@@ -32,13 +32,13 @@
 ;; The last splitpoint uses +nan.0 for pt and represents the "else"
 (struct sp (cidx bexpr point) #:prefab)
 
-(define (combine-alts best-option start-prog ctx)
+(define (combine-alts best-option start-prog ctx pcontext)
   (match-define (option splitindices alts pts expr _) best-option)
   (match splitindices
     [(list (si cidx _)) (list-ref alts cidx)]
     [_
      (timeline-event! 'bsearch)
-     (define splitpoints (sindices->spoints pts expr alts splitindices start-prog ctx))
+     (define splitpoints (sindices->spoints pts expr alts splitindices start-prog ctx pcontext))
 
      (define expr*
        (for/fold ([expr (alt-expr (list-ref alts (sp-cidx (last splitpoints))))])
@@ -119,8 +119,8 @@
 ;; float form always come from the range [f(idx1), f(idx2)). If the
 ;; float form of a split is f(idx2), or entirely outside that range,
 ;; problems may arise.
-(define/contract (sindices->spoints points expr alts sindices start-prog ctx)
-  (-> (listof vector?) any/c (listof alt?) (listof si?) any/c context? valid-splitpoints?)
+(define/contract (sindices->spoints points expr alts sindices start-prog ctx pcontext)
+  (-> (listof vector?) any/c (listof alt?) (listof si?) any/c context? pcontext? valid-splitpoints?)
   (define repr (repr-of expr ctx))
 
   (define eval-expr (compile-prog expr ctx))
@@ -135,7 +135,7 @@
     (and start-prog (make-real-compiler (list (prog->spec start-prog)) (list ctx*))))
 
   (define (prepend-macro v)
-    (prepend-argument start-real-compiler v (*pcontext*)))
+    (prepend-argument start-real-compiler v pcontext))
 
   (define (find-split expr1 expr2 v1 v2)
     (define (pred v)

src/core/explain.rkt

@@ -37,8 +37,8 @@
 
 (define (actual-errors expr ctx pcontext)
   (match-define (cons subexprs pt-errorss)
-    (parameterize ([*pcontext* pcontext])
-      (flip-lists (hash->list (first (compute-local-errors (list (all-subexpressions expr)) ctx))))))
+    (flip-lists
+     (hash->list (first (compute-local-errors (list (all-subexpressions expr)) ctx pcontext)))))
 
   (define pt-worst-subexpr
     (append* (reap [sow]

src/core/localize.rkt

@@ -58,10 +58,10 @@
 (define (make-matrix roots pcontext)
   (for/vector #:length (vector-length roots)
               ([node (in-vector roots)])
-    (make-vector (pcontext-length (*pcontext*)))))
+    (make-vector (pcontext-length pcontext))))
 
 ; Compute local error or each sampled point at each node in `prog`.
-(define (compute-local-errors subexprss ctx)
+(define (compute-local-errors subexprss ctx pcontext)
   (define exprs-list (append* subexprss)) ; unroll subexprss
   (define reprs-list (map (curryr repr-of ctx) exprs-list))
   (define ctx-list
@@ -75,9 +75,9 @@
 
   (define subexprs-fn (eval-progs-real (map prog->spec exprs-list) ctx-list))
 
-  (define errs (make-matrix roots (*pcontext*)))
+  (define errs (make-matrix roots pcontext))
 
-  (for ([(pt ex) (in-pcontext (*pcontext*))]
+  (for ([(pt ex) (in-pcontext pcontext)]
         [pt-idx (in-naturals)])
     (define exacts (list->vector (subexprs-fn pt)))
     (define (get-exact idx)
@@ -110,7 +110,7 @@
         ((representation-bf->repr repr) 0.bf))))
 
 ;; Compute local error or each sampled point at each node in `prog`.
-(define (compute-errors subexprss ctx)
+(define (compute-errors subexprss ctx pcontext)
   ;; We compute the actual (float) result
   (define exprs-list (append* subexprss)) ; unroll subexprss
   (define actual-value-fn (compile-progs exprs-list ctx))
@@ -146,14 +146,14 @@
   (define nodes (batch-nodes expr-batch))
   (define roots (batch-roots expr-batch))
 
-  (define ulp-errs (make-matrix roots (*pcontext*)))
-  (define exacts-out (make-matrix roots (*pcontext*)))
-  (define approx-out (make-matrix roots (*pcontext*)))
-  (define true-error-out (make-matrix roots (*pcontext*)))
+  (define ulp-errs (make-matrix roots pcontext))
+  (define exacts-out (make-matrix roots pcontext))
+  (define approx-out (make-matrix roots pcontext))
+  (define true-error-out (make-matrix roots pcontext))
 
   (define spec-vec (list->vector spec-list))
   (define ctx-vec (list->vector ctx-list))
-  (for ([(pt ex) (in-pcontext (*pcontext*))]
+  (for ([(pt ex) (in-pcontext pcontext)]
         [pt-idx (in-naturals)])
 
     (define exacts (list->vector (subexprs-fn pt)))
@@ -211,8 +211,8 @@
 ;; Compute the local error of every subexpression of `prog`
 ;; and returns the error information as an S-expr in the
 ;; same shape as `prog`
-(define (local-error-as-tree expr ctx)
-  (define data-hash (first (compute-errors (list (all-subexpressions expr)) ctx)))
+(define (local-error-as-tree expr ctx pcontext)
+  (define data-hash (first (compute-errors (list (all-subexpressions expr)) ctx pcontext)))
 
   (define (translate-booleans value)
     (match value

src/core/mainloop.rkt

@@ -32,6 +32,7 @@
 
 ;; Starting program for the current run
 (define *start-prog* (make-parameter #f))
+(define *pcontext* (make-parameter #f))
 
 ;; These high-level functions give the high-level workflow of Herbie:
 ;; - Initial steps: explain, preprocessing, initialize the alt table
@@ -251,9 +252,10 @@
           (equal? (representation-type repr) 'real)
           (not (null? (context-vars ctx)))
           (get-fpcore-impl '<= '() (list repr repr)))
-     (define opts (pareto-regimes (sort alts < #:key (curryr alt-cost repr)) start-prog ctx))
+     (define opts
+       (pareto-regimes (sort alts < #:key (curryr alt-cost repr)) start-prog ctx (*pcontext*)))
      (for/list ([opt (in-list opts)])
-       (combine-alts opt start-prog ctx))]
+       (combine-alts opt start-prog ctx (*pcontext*)))]
     [else (list (argmin score-alt alts))]))
 
 (define (add-derivations! alts)

src/core/points.rkt

@@ -5,8 +5,7 @@
          "batch.rkt"
          "compiler.rkt")
 
-(provide *pcontext*
-         in-pcontext
+(provide in-pcontext
          mk-pcontext
          for/pcontext
          pcontext?
@@ -21,7 +20,6 @@
 ;; ground-truth information. They contain 1) a set of sampled input
 ;; points; and 2) a ground-truth output for each input.
 
-(define *pcontext* (make-parameter #f))
 (struct pcontext (points exacts) #:prefab)
 
 (define (in-pcontext pcontext)

src/core/regimes.rkt

@@ -25,9 +25,9 @@
   [(define (write-proc opt port mode)
      (fprintf port "#<option ~a>" (option-split-indices opt)))])
 
-(define (pareto-regimes sorted start-prog ctx)
+(define (pareto-regimes sorted start-prog ctx pcontext)
   (timeline-event! 'regimes)
-  (define err-lsts (batch-errors (map alt-expr sorted) (*pcontext*) ctx))
+  (define err-lsts (batch-errors (map alt-expr sorted) pcontext ctx))
   (define branches
     (if (null? sorted)
         '()
@@ -44,14 +44,15 @@
       ; Only return one option if not pareto mode
       [(and (not (*pareto-mode*)) (not (equal? alts sorted))) '()]
       [else
-       (define-values (opt new-errs) (infer-splitpoints branch-exprs alts err-lsts #:errs errs ctx))
+       (define-values (opt new-errs)
+         (infer-splitpoints branch-exprs alts err-lsts #:errs errs ctx pcontext))
        (define high (si-cidx (argmax (λ (x) (si-cidx x)) (option-split-indices opt))))
        (cons opt (loop (take alts high) new-errs (take err-lsts high)))])))
 
 ;; `infer-splitpoints` and `combine-alts` are split so the mainloop
 ;; can insert a timeline break between them.
 
-(define (infer-splitpoints branch-exprs alts err-lsts* #:errs [cerrs (hash)] ctx)
+(define (infer-splitpoints branch-exprs alts err-lsts* #:errs [cerrs (hash)] ctx pcontext)
   (timeline-push! 'inputs (map (compose ~a alt-expr) alts))
   (define sorted-bexprs
     (sort branch-exprs (lambda (x y) (< (hash-ref cerrs x -1) (hash-ref cerrs y -1)))))
@@ -67,7 +68,7 @@
               ([bexpr sorted-bexprs]
                ;; stop if we've computed this (and following) branch-expr on more alts and it's still worse
                #:break (> (hash-ref cerrs bexpr -1) best-err))
-      (define opt (option-on-expr alts err-lsts bexpr ctx))
+      (define opt (option-on-expr alts err-lsts bexpr ctx pcontext))
       (define err
         (+ (errors-score (option-errors opt))
            (length (option-split-indices opt)))) ;; one-bit penalty per split
@@ -113,14 +114,14 @@
              #:when (critical-subexpression? expr subexpr))
     subexpr))
 
-(define (option-on-expr alts err-lsts expr ctx)
+(define (option-on-expr alts err-lsts expr ctx pcontext)
   (define timeline-stop! (timeline-start! 'times (~a expr)))
 
   (define fn (compile-prog expr ctx))
   (define repr (repr-of expr ctx))
 
   (define big-table ; pt ; splitval ; alt1-err ; alt2-err ; ...
-    (for/list ([(pt ex) (in-pcontext (*pcontext*))]
+    (for/list ([(pt ex) (in-pcontext pcontext)]
                [err-lst err-lsts])
       (list* (fn pt) pt err-lst)))
   (match-define (list splitvals* pts* err-lsts* ...)
@@ -151,27 +152,27 @@
 
 (module+ test
   (define ctx (make-debug-context '(x)))
-  (parameterize ([*pcontext* (mk-pcontext '(#(0.5) #(4.0)) '(1.0 1.0))])
-    (define alts (map make-alt (list '(fmin.f64 x 1) '(fmax.f64 x 1))))
-    (define err-lsts `((,(expt 2.0 53) 1.0) (1.0 ,(expt 2.0 53))))
-
-    (define (test-regimes expr goal)
-      (check (lambda (x y) (equal? (map si-cidx (option-split-indices x)) y))
-             (option-on-expr alts err-lsts expr ctx)
-             goal))
-
-    ;; This is a basic sanity test
-    (test-regimes 'x '(1 0))
-
-    ;; This test ensures we handle equal points correctly. All points
-    ;; are equal along the `1` axis, so we should only get one
-    ;; splitpoint (the second, since it is better at the further point).
-    (test-regimes (literal 1 'binary64) '(0))
-
-    (test-regimes `(if (==.f64 x ,(literal 0.5 'binary64))
-                       ,(literal 1 'binary64)
-                       (NAN.f64))
-                  '(1 0))))
+  (define pctx (mk-pcontext '(#(0.5) #(4.0)) '(1.0 1.0)))
+  (define alts (map make-alt (list '(fmin.f64 x 1) '(fmax.f64 x 1))))
+  (define err-lsts `((,(expt 2.0 53) 1.0) (1.0 ,(expt 2.0 53))))
+
+  (define (test-regimes expr goal)
+    (check (lambda (x y) (equal? (map si-cidx (option-split-indices x)) y))
+           (option-on-expr alts err-lsts expr ctx pctx)
+           goal))
+
+  ;; This is a basic sanity test
+  (test-regimes 'x '(1 0))
+
+  ;; This test ensures we handle equal points correctly. All points
+  ;; are equal along the `1` axis, so we should only get one
+  ;; splitpoint (the second, since it is better at the further point).
+  (test-regimes (literal 1 'binary64) '(0))
+
+  (test-regimes `(if (==.f64 x ,(literal 0.5 'binary64))
+                     ,(literal 1 'binary64)
+                     (NAN.f64))
+                '(1 0)))
 
 ;; Given error-lsts, returns a list of sp objects representing where the optimal splitpoints are.
 (define (valid-splitindices? can-split? split-indices)