Merge pull request #1238 from herbie-fp/compile-preprocessing

Compile away preprocessing

Author: pavpanchekha

src/core/mainloop.rkt

@@ -50,10 +50,17 @@
   (define alternatives (extract!))
   (timeline-event! 'preprocess)
   (for/list ([altn alternatives])
-    (define expr (alt-expr altn))
-    (define preprocessing (alt-preprocessing altn))
-    (alt-add-preprocessing altn
-                           (remove-unnecessary-preprocessing expr context pcontext preprocessing))))
+    (apply-preprocessing altn context pcontext)))
+
+(define (apply-preprocessing altn context pcontext)
+  (define expr (alt-expr altn))
+  (define initial-preprocessing (alt-preprocessing altn))
+  (define useful-preprocessing
+    (remove-unnecessary-preprocessing expr context pcontext initial-preprocessing))
+  (define expr*
+    (for/fold ([expr expr]) ([preprocessing (in-list (reverse useful-preprocessing))])
+      (compile-preprocessing expr context preprocessing)))
+  (alt expr* 'add-preprocessing (list altn) '()))
 
 (define (extract!)
   (timeline-push-alts! '())

src/core/preprocess.rkt

@@ -1,5 +1,6 @@
 #lang racket
 
+(require math/bigfloat)
 (require "../syntax/platform.rkt"
          "../syntax/syntax.rkt"
          "../syntax/types.rkt"
@@ -17,37 +18,46 @@
 
 (provide find-preprocessing
          preprocess-pcontext
-         remove-unnecessary-preprocessing)
+         remove-unnecessary-preprocessing
+         compile-preprocessing)
 
-(define (has-fabs-neg-impls? repr)
-  (and (get-fpcore-impl '- (repr->prop repr) (list repr))
-       (get-fpcore-impl 'fabs (repr->prop repr) (list repr))))
+(define (has-fabs-impl? repr)
+  (get-fpcore-impl 'fabs (repr->prop repr) (list repr)))
+
+(define (has-fmin-fmax-impl? repr)
+  (and (get-fpcore-impl 'fmin (repr->prop repr) (list repr repr))
+       (get-fpcore-impl 'fmax (repr->prop repr) (list repr repr))))
 
 (define (has-copysign-impl? repr)
-  (get-fpcore-impl 'copysign (repr->prop repr) (list repr repr)))
+  (and (get-fpcore-impl '* (repr->prop repr) (list repr repr))
+       (get-fpcore-impl 'copysign (repr->prop repr) (list repr repr))))
 
 ;; The even identities: f(x) = f(-x)
 ;; Requires `neg` and `fabs` operator implementations.
 (define (make-even-identities spec ctx)
   (for/list ([var (in-list (context-vars ctx))]
              [repr (in-list (context-var-reprs ctx))]
-             #:when (has-fabs-neg-impls? repr))
+             #:when (has-fabs-impl? repr))
     (cons `(abs ,var) (replace-expression spec var `(neg ,var)))))
 
 ;; The odd identities: f(x) = -f(-x)
 ;; Requires `neg` and `fabs` operator implementations.
 (define (make-odd-identities spec ctx)
   (for/list ([var (in-list (context-vars ctx))]
              [repr (in-list (context-var-reprs ctx))]
-             #:when (and (has-fabs-neg-impls? repr) (has-copysign-impl? repr)))
+             #:when (and (has-fabs-impl? repr) (has-copysign-impl? (context-repr ctx))))
     (cons `(negabs ,var) (replace-expression `(neg ,spec) var `(neg ,var)))))
 
-;; Swap identities: f(a, b) = f(b, a)
-(define (make-swap-identities spec ctx)
+;; Sort identities: f(a, b) = f(b, a)
+;; TODO: require both vars have the same repr
+(define (make-sort-identities spec ctx)
   (define pairs (combinations (context-vars ctx) 2))
-  (for/list ([pair (in-list pairs)])
+  (for/list ([pair (in-list pairs)]
+             ;; Can only sort same-repr variables
+             #:when (equal? (context-lookup ctx (first pair)) (context-lookup ctx (second pair)))
+             #:when (has-fmin-fmax-impl? (context-lookup ctx (first pair))))
     (match-define (list a b) pair)
-    (cons `(swap ,a ,b) (replace-vars `((,a . ,b) (,b . ,a)) spec))))
+    (cons `(sort ,a ,b) (replace-vars `((,a . ,b) (,b . ,a)) spec))))
 
 ;; See https://pavpanchekha.com/blog/symmetric-expressions.html
 (define (find-preprocessing expr ctx)
@@ -56,8 +66,8 @@
   ;; identities
   (define even-identities (make-even-identities spec ctx))
   (define odd-identities (make-odd-identities spec ctx))
-  (define swap-identities (make-swap-identities spec ctx))
-  (define identities (append even-identities odd-identities swap-identities))
+  (define sort-identities (make-sort-identities spec ctx))
+  (define identities (append even-identities odd-identities sort-identities))
 
   ;; make egg runner
   (define rules (*sound-rules*))
@@ -69,8 +79,6 @@
                  (make-list (vector-length (batch-roots batch)) (context-repr ctx))
                  `((,rules . ((node . ,(*node-limit*)))))))
 
-  ;; TODO : FIGURE HOW TO IMPLEMENT PREPROCESS
-
   ;; collect equalities
   (define abs-instrs
     (for/list ([(ident spec*) (in-dict even-identities)]
@@ -82,28 +90,12 @@
                #:when (egraph-equal? runner spec spec*))
       ident))
 
-  (define swaps
-    (for/list ([(ident spec*) (in-dict swap-identities)]
-               #:when (egraph-equal? runner spec spec*))
-      (match-define (list 'swap a b) ident)
-      (list a b)))
-  (define components (connected-components (context-vars ctx) swaps))
   (define sort-instrs
-    (for/list ([component (in-list components)]
-               #:when (> (length component) 1))
-      (cons 'sort component)))
-
+    (for/list ([(ident spec*) (in-dict sort-identities)]
+               #:when (egraph-equal? runner spec spec*))
+      ident))
   (append abs-instrs negabs-instrs sort-instrs))
 
-(define (connected-components variables swaps)
-  (define components (disjoint-set (length variables)))
-  (for ([swap (in-list swaps)])
-    (match-define (list a b) swap)
-    (disjoint-set-union! components
-                         (disjoint-set-find! components (index-of variables a))
-                         (disjoint-set-find! components (index-of variables b))))
-  (group-by (compose (curry disjoint-set-find! components) (curry index-of variables)) variables))
-
 (define (preprocess-pcontext context pcontext preprocessing)
   (define preprocess
     (apply compose
@@ -130,30 +122,28 @@
   (define variables (context-vars context))
   (define sort* (curryr sort (curryr </total (context-repr context))))
   (match instruction
-    [(list 'sort component ...)
-     (define indices (indexes-where variables (curryr member component)))
+    [(list 'sort a b)
+     (define indices (indexes-where variables (curry set-member? (list a b))))
+     (define repr (context-lookup context a))
      (lambda (x y)
        (define subsequence (map (curry vector-ref x) indices))
-       (define sorted (sort* subsequence))
+       (define sorted (sort subsequence (curryr </total repr)))
        (values (vector-set* x indices sorted) y))]
     [(list 'abs variable)
      (define index (index-of variables variable))
      (define var-repr (context-lookup context variable))
-     (define abs-proc (impl-info (get-fpcore-impl 'fabs (repr->prop var-repr) (list var-repr)) 'fl))
-     (lambda (x y) (values (vector-update x index abs-proc) y))]
+     (define fabs (impl-info (get-fpcore-impl 'fabs (repr->prop var-repr) (list var-repr)) 'fl))
+     (lambda (x y) (values (vector-update x index fabs) y))]
     [(list 'negabs variable)
      (define index (index-of variables variable))
      (define var-repr (context-lookup context variable))
-     (define neg-var (impl-info (get-fpcore-impl '- (repr->prop var-repr) (list var-repr)) 'fl))
-
      (define repr (context-repr context))
-     (define neg-expr (impl-info (get-fpcore-impl '- (repr->prop repr) (list repr)) 'fl))
-
+     (define fabs (impl-info (get-fpcore-impl 'fabs (repr->prop var-repr) (list var-repr)) 'fl))
+     (define mul (impl-info (get-fpcore-impl '* (repr->prop repr) (list repr repr)) 'fl))
+     (define copysign (impl-info (get-fpcore-impl 'copysign (repr->prop repr) (list repr repr)) 'fl))
+     (define repr1 ((representation-bf->repr repr) 1.bf))
      (lambda (x y)
-       ;; Negation is involutive, i.e. it is its own inverse, so t^1(y') = -y'
-       (if (negative? (repr->real (vector-ref x index) (context-repr context)))
-           (values (vector-update x index neg-var) (neg-expr y))
-           (values x y)))]))
+       (values (vector-update x index fabs) (mul (copysign repr1 (vector-ref x index)) y)))]))
 
 ; until fixed point, iterate through preprocessing attempting to drop preprocessing with no effect on error
 (define (remove-unnecessary-preprocessing expression
@@ -182,3 +172,25 @@
   (define pcontext2 (preprocess-pcontext context pcontext preprocessing2))
   (<= (errors-score (errors expression pcontext1 context))
       (errors-score (errors expression pcontext2 context))))
+
+(define (compile-preprocessing expression context preprocessing)
+  (match preprocessing
+    ; Not handled yet
+    [(list 'sort a b)
+     (define repr (context-lookup context a))
+     (define fmin (get-fpcore-impl 'fmin (repr->prop repr) (list repr repr)))
+     (define fmax (get-fpcore-impl 'fmax (repr->prop repr) (list repr repr)))
+     (replace-vars (list (cons a `(,fmin ,a ,b)) (cons b `(,fmax ,a ,b))) expression)]
+    [(list 'abs var)
+     (define repr (context-lookup context var))
+     (define fabs (get-fpcore-impl 'fabs (repr->prop repr) (list repr)))
+     (define replacement `(,fabs ,var))
+     (replace-expression expression var replacement)]
+    [(list 'negabs var)
+     (define repr (context-lookup context var))
+     (define fabs (get-fpcore-impl 'fabs (repr->prop repr) (list repr)))
+     (define replacement `(,fabs ,var))
+     (define mul (get-fpcore-impl '* (repr->prop repr) (list repr repr)))
+     (define copysign (get-fpcore-impl 'copysign (repr->prop repr) (list repr repr)))
+     `(,mul (,copysign ,(literal 1 (representation-name repr)) ,var)
+            ,(replace-expression expression var replacement))]))

src/core/rules.rkt

@@ -241,6 +241,15 @@
   [fabs-cbrt (fabs (/ (cbrt a) a)) (/ (cbrt a) a)]
   [fabs-cbrt-rev (/ (cbrt a) a) (fabs (/ (cbrt a) a))])
 
+; Copysign
+(define-rules arithmetic
+  [copysign-neg (copysign a (neg b)) (neg (copysign a b))]
+  [neg-copysign (neg (copysign a b)) (copysign a (neg b))]
+  [copysign-other-neg (copysign (neg a) b) (copysign a b)]
+  [copysign-fabs (copysign a (fabs b)) (fabs a)]
+  [copysign-other-fabs (copysign (fabs a) b) (copysign a b)]
+  [fabs-copysign (fabs (copysign a b)) (fabs a)])
+
 ; Square root
 (define-rules arithmetic
   [sqrt-pow2 (pow (sqrt x) y) (pow x (/ y 2))]
@@ -284,6 +293,11 @@
   [cbrt-div-cbrt (/ (cbrt x) (fabs (cbrt x))) (copysign 1 x)]
   [cbrt-div-cbrt2 (/ (fabs (cbrt x)) (cbrt x)) (copysign 1 x)])
 
+; Min and max
+(define-rules arithmetic
+  [fmin-swap (fmin a b) (fmin b a)]
+  [fmax-swap (fmax a b) (fmax b a)])
+
 ; Exponentials
 (define-rules exponents
   [add-log-exp x (log (exp x))]

src/reports/common.rkt

@@ -155,117 +155,15 @@
       (core->tex prog* #:loc (and loc (cons 2 loc)) #:color "blue")
       "ERROR"))
 
-(define (combine-fpcore-instruction i e c)
-  (match i
-    [(list 'abs x)
-     (define x* (string->symbol (string-append (symbol->string x) "_m")))
-     (define e* (replace-expression e x x*))
-     (define p (index-of (context-vars c) x))
-     (define c* (struct-copy context c [vars (list-set (context-vars c) p x*)]))
-     (cons e* c*)]
-    [(list 'negabs x)
-     (define x-string (symbol->string x))
-     (define x-sign (string->symbol (string-append x-string "_s")))
-     (define x* (string->symbol (string-append x-string "_m")))
-     (define p (index-of (context-vars c) x))
-     (define r (list-ref (context-var-reprs c) p))
-     (define c* (struct-copy context c [vars (list-set (context-vars c) p x*)]))
-     (define c** (context-extend c* x-sign r))
-     (define *-impl (get-fpcore-impl '* (repr->prop (context-repr c)) (list r (context-repr c))))
-     (define e* (list *-impl x-sign (replace-expression e x x*)))
-     (cons e* c**)]
-    [_ (cons e c)]))
-
-(define (format-prelude-instruction instruction ctx ctx* language converter)
-  (define (converter* e c)
-    (define fpcore (program->fpcore e c))
-    (define output (converter fpcore "code"))
-    (define lines (string-split output "\n"))
-    (match language
-      ["FPCore" (pretty-format e #:mode 'display)]
-      ["Fortran" (string-trim (third lines) #px"\\s+code\\s+=\\s+")]
-      ["MATLAB" (string-trim (second lines) #px"\\s+tmp\\s+=\\s+")]
-      ["Wolfram" (string-trim (first lines) #px".*:=\\s+")]
-      ["TeX" output]
-      [_ (string-trim (second lines) #px"\\s+return\\s+")]))
-  (match instruction
-    [(list 'abs x)
-     (define x* (string->symbol (string-append (symbol->string x) "_m")))
-     (define r (list-ref (context-var-reprs ctx) (index-of (context-vars ctx) x)))
-     (define fabs-impl (get-fpcore-impl 'fabs (repr->prop r) (list r)))
-     (define e (list fabs-impl x))
-     (define c (context (list x) r r))
-     (list (format "~a = ~a" x* (converter* e c)))]
-    [(list 'negabs x)
-     ; TODO: why are x* and x-sign unused?
-     (define x* (string->symbol (format "~a_m" x)))
-     (define r (context-lookup ctx x))
-     (define fabs-impl (get-fpcore-impl 'fabs (repr->prop r) (list r)))
-     (define copysign-impl (get-fpcore-impl 'copysign (repr->prop r) (list r r)))
-     (define e* (list fabs-impl x))
-     (define x-sign (string->symbol (format "~a_s" x)))
-     (define e-sign (list copysign-impl (literal 1 (representation-name r)) x))
-     (define c (context (list x) r r))
-     (list (format "~a = ~a" (format "~a\\_m" x) (converter* e* c))
-           (format "~a = ~a" (format "~a\\_s" x) (converter* e-sign c)))]
-    [(list 'sort vs ...)
-     (define vs (context-vars ctx))
-     (define vs* (context-vars ctx*))
-     ;; We added some sign-* variables to the front of the variable
-     ;; list in `ctx*`, we only want the originals here
-     (list (format-sort-instruction (take-right vs* (length vs)) language))]))
-
-(define (format-sort-instruction vs l)
-  (match l
-    ["C" (format "assert(~a);" (format-less-than-condition vs))]
-    ["Java" (format "assert ~a;" (format-less-than-condition vs))]
-    ["Python"
-     (define comma-joined (comma-join vs))
-     (format "[~a] = sort([~a])" comma-joined comma-joined)]
-    ["Julia"
-     (define comma-joined (comma-join vs))
-     (format "~a = sort([~a])" comma-joined comma-joined)]
-    ["MATLAB"
-     (define comma-joined (comma-join vs))
-     (format "~a = num2cell(sort([~a])){:}" comma-joined comma-joined)]
-    ["TeX"
-     (define comma-joined (comma-join vs))
-     (format "[~a] = \\mathsf{sort}([~a])\\\\" comma-joined comma-joined)]
-    [_
-     (match vs
-       [(list x y) (format sort-note (format "~a and ~a" x y))]
-       [(list vs ...)
-        (format sort-note
-                (string-join (map ~a vs)
-                             ", "
-                             ;; "Lil Jon, he always tells the truth"
-                             #:before-last ", and "))])]))
-
-(define (format-less-than-condition variables)
-  (string-join (for/list ([a (in-list variables)]
-                          [b (in-list (cdr variables))])
-                 (format "~a < ~a" a b))
-               " && "))
-
-(define (comma-join vs)
-  (string-join (map ~a vs) ", "))
-
-(define sort-note "NOTE: ~a should be sorted in increasing order before calling this function.")
-
 (define (render-program expr
                         ctx
                         #:ident [identifier #f]
                         #:pre [precondition '(TRUE)]
                         #:instructions [instructions empty])
   (define output-repr (context-repr ctx))
-  (match-define (cons expr* ctx*)
-    (foldl (match-lambda*
-             [(list i (cons e c)) (combine-fpcore-instruction i e c)])
-           (cons expr ctx)
-           instructions))
   (define out-prog
     (parameterize ([*expr-cse-able?* at-least-two-ops?])
-      (core-cse (program->fpcore expr* ctx* #:ident identifier))))
+      (core-cse (program->fpcore expr ctx #:ident identifier))))
 
   (define output-prec (representation-name output-repr))
   (define out-prog* (fpcore-add-props out-prog (list ':precision output-prec)))
@@ -281,19 +179,7 @@
                     (symbol->string identifier)
                     "code"))
               (define out (converter out-prog* name))
-              (define prelude-lines
-                (string-join
-                 (append-map (lambda (instruction)
-                               (format-prelude-instruction instruction ctx ctx* lang converter))
-                             instructions)
-                 (if (equal? lang "TeX") "\\\\\n" "\n")
-                 #:after-last "\n"))
-              (sow (cons lang
-                         ((if (equal? lang "TeX")
-                              (curry format "\\begin{array}{l}\n~a\\\\\n~a\\end{array}\n")
-                              string-append)
-                          prelude-lines
-                          out)))))))
+              (sow (cons lang out))))))
 
   (define math-out (dict-ref versions "TeX" ""))