Merge pull request #1127 from herbie-fp/simplify-patch

Add `(hole repr spec)` terms to simplify Taylor expansion

Author: pavpanchekha

src/core/batch.rkt

@@ -26,6 +26,7 @@
 (define (expr-recurse expr f)
   (match expr
     [(approx spec impl) (approx (f spec) (f impl))]
+    [(hole precision spec) (hole precision (f spec))]
     [(list op args ...) (cons op (map f args))]
     [_ expr]))
 

src/core/derivations.rkt

@@ -35,35 +35,13 @@
   (define (build! altn)
     (match altn
       ; recursive rewrite using egg (impl -> impl)
-      [(alt expr `(rr ,loc ,(? egg-runner? runner) #f) `(,prev) _)
+      [(alt expr `(,(or 'rr 'simplify) ,loc ,(? egg-runner? runner) #f) `(,prev) _)
        (define start-expr (location-get loc (alt-expr prev)))
        (define end-expr (location-get loc expr))
        (define rewrite (cons start-expr end-expr))
        (hash-set! alt->query&rws (altn->key altn) (cons runner rewrite))
        (hash-update! query->rws runner (lambda (rws) (set-add rws rewrite)) '())]
 
-      ; simplify using egg
-      ;  usually: impl -> impl
-      ;  taylor: spec -> approx (_, impl)
-      [(alt expr `(simplify ,loc ,(? egg-runner? runner) #f) `(,prev) _)
-       (define rewrite
-         (match (alt-event prev)
-           [(list 'taylor _ ...)
-            ; simplify after taylor: spec -> approx (_, impl)
-            (define start-expr (location-get loc (alt-expr prev)))
-            (define end-expr (location-get loc expr))
-            (unless (approx? end-expr)
-              (error 'make-proof-tables "expected approx node, got ~a" end-expr))
-            (cons start-expr (approx-impl end-expr))]
-           [_
-            ; simplify after other: impl -> impl
-            (define start-expr (location-get loc (alt-expr prev)))
-            (define end-expr (location-get loc expr))
-            (cons start-expr end-expr)]))
-
-       (hash-set! alt->query&rws (altn->key altn) (cons runner rewrite))
-       (hash-update! query->rws runner (lambda (rws) (set-add rws rewrite)) '())]
-
       ; everything else
       [_ (void)])
 

src/core/egg-herbie.rkt

@@ -132,22 +132,23 @@
   (for ([node (in-vector (batch-nodes insert-batch))]
         [root? (in-vector root-mask)]
         [n (in-naturals)])
-    (define node*
+    (define idx
       (match node
-        [(literal v _) v]
-        [(? number?) node]
-        [(? symbol?) (normalize-var node)]
+        [(literal v _) (insert-node! v root?)]
+        [(? number?) (insert-node! node root?)]
+        [(? symbol?) (insert-node! (normalize-var node) root?)]
+        [(hole prec spec) (remap spec)] ; "hole" terms currently disappear
         [(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)]))
+         (insert-node! (list '$approx (remap spec) (remap impl)) root?)]
+        [(list op (app remap args) ...) (insert-node! (cons op args) root?)]))
 
-    (vector-set! mappings n (insert-node! node* root?)))
+    (vector-set! mappings n idx))
 
   (for/list ([root (in-vector (batch-roots insert-batch))])
     (remap root)))
@@ -276,6 +277,7 @@
                     (hash-set! egg->herbie-dict replacement (cons expr (context-lookup ctx expr)))
                     replacement))]
       [(approx spec impl) (list '$approx (loop spec) (loop impl))]
+      [(hole precision spec) (loop spec)]
       [(list op args ...) (cons op (map loop args))])))
 
 (define (flatten-let expr)

src/core/patch.rkt

@@ -55,13 +55,8 @@
                 [outputs (in-list simplification-options)])
             (match-define (cons _ simplified) outputs)
             (define prev (car (alt-prevs altn)))
-            (for ([batchreff (in-list simplified)])
-              (define spec (prog->spec (debatchref (alt-expr prev))))
-              (define idx ; Munge
-                (mutable-batch-push! global-batch-mutable
-                                     (approx (mutable-batch-munge! global-batch-mutable spec)
-                                             (batchref-idx batchreff))))
-              (sow (alt (batchref global-batch idx) `(simplify ,runner #f) (list altn) '()))))
+            (for ([bref (in-list simplified)])
+              (sow (alt bref `(simplify ,runner #f) (list altn) '()))))
           (batch-copy-mutable-nodes! global-batch global-batch-mutable))) ; Update global-batch
 
   (timeline-push! 'count (length approxs) (length simplified))
@@ -80,25 +75,26 @@
                               #;(log ,log-x ,exp-x))))
 
 (define (taylor-alts starting-exprs altns global-batch)
-  (define exprs
-    (for/list ([expr (in-list starting-exprs)])
-      (prog->spec expr)))
-  (define free-vars (map free-variables exprs))
+  (define specs (map prog->spec starting-exprs))
+  (define free-vars (map free-variables specs))
   (define vars (context-vars (*context*)))
 
   (reap [sow]
         (define global-batch-mutable (batch->mutable-batch global-batch)) ; Create a mutable batch
         (for* ([var (in-list vars)]
                [transform-type transforms-to-try])
           (match-define (list name f finv) transform-type)
-          (define timeline-stop! (timeline-start! 'series (~a exprs) (~a var) (~a name)))
-          (define genexprs (approximate exprs var #:transform (cons f finv)))
+          (define timeline-stop! (timeline-start! 'series (~a specs) (~a var) (~a name)))
+          (define genexprs (approximate specs var #:transform (cons f finv)))
           (for ([genexpr (in-list genexprs)]
+                [spec (in-list specs)]
+                [expr (in-list starting-exprs)]
                 [altn (in-list altns)]
                 [fv (in-list free-vars)]
-                #:when (member var fv)) ; check whether var exists in expr at all
+                #:when (set-member? fv var)) ; check whether var exists in expr at all
             (for ([i (in-range (*taylor-order-limit*))])
-              (define gen (genexpr))
+              (define repr (repr-of expr (*context*)))
+              (define gen (approx spec (hole (representation-name repr) (genexpr))))
               (define idx (mutable-batch-munge! global-batch-mutable gen)) ; Munge gen
               (sow (alt (batchref global-batch idx) `(taylor ,name ,var) (list altn) '()))))
           (timeline-stop!))

src/core/programs.rkt

@@ -40,6 +40,7 @@
     [(literal val precision) (get-representation precision)]
     [(? variable?) (context-lookup ctx node)]
     [(approx _ impl) (repr-of-node batch impl ctx)]
+    [(hole precision spec) (get-representation precision)]
     [(list 'if cond ift iff) (repr-of-node batch ift ctx)]
     [(list op args ...) (impl-info op 'otype)]))
 

src/reports/history.rkt

@@ -73,16 +73,17 @@
   (values (format-accuracy err repr-bits #:unit "%")
           (format "~a on training set" (format-accuracy err2 repr-bits #:unit "%"))))
 
-(define (remove-literals expr)
-  (match expr
-    [(? symbol?) expr]
-    [(? number?) expr]
-    [(? literal?) (literal-value expr)]
-    [(approx spec impl) (approx (remove-literals spec) (remove-literals impl))]
-    [(list op args ...) (cons op (map remove-literals args))]))
-
 (define (expr->fpcore expr ctx #:ident [ident #f])
-  (list 'FPCore (context-vars ctx) (remove-literals expr)))
+  (list 'FPCore
+        (context-vars ctx)
+        (let loop ([expr expr])
+          (match expr
+            [(? symbol?) expr]
+            [(? number?) expr]
+            [(? literal?) (literal-value expr)]
+            [(approx spec impl) (loop impl)]
+            [(hole precision spec) (loop spec)]
+            [(list op args ...) (cons op (map loop args))]))))
 
 (define (mixed->fpcore expr ctx)
   (define expr*
@@ -92,6 +93,7 @@
         [(? number?) expr]
         [(? literal?) (literal-value expr)]
         [(approx _ impl) (loop impl)]
+        [(hole precision spec) (loop spec)]
         [`(if ,cond ,ift ,iff)
          `(if ,(loop cond)
               ,(loop ift)

src/syntax/syntax.rkt

@@ -10,6 +10,7 @@
 
 (provide (struct-out literal)
          (struct-out approx)
+         (struct-out hole)
          variable?
          constant-operator?
          operator-exists?
@@ -493,9 +494,12 @@
 (struct literal (value precision) #:prefab)
 
 ;; An approximation of a specification by
-;; an arbitrary floating-point expression.
+;; a floating-point expression.
 (struct approx (spec impl) #:prefab)
 
+;; An unknown floating-point expression that implements a given spec
+(struct hole (precision spec) #:prefab)
+
 ;; name -> (vars repr body)	;; name -> (vars prec body)
 (define *functions* (make-parameter (make-hasheq)))