Merge pull request #1325 from herbie-fp/codex/plan-changes-to-remove-if-costs

Normalize if operator

Author: pavpanchekha

infra/softposit.rkt

@@ -82,7 +82,6 @@
 
 ;;;;;;;;;;;;;;;;;;;;;;;;;;;;; EMPTY PLATFORM ;;;;;;;;;;;;;;;;;;;;;;;;
 
-(define-if #:cost 1)
 
 ;;;;;;;;;;;;;;;;;;;;;;;;;;;;; REPRESENTATIONS ;;;;;;;;;;;;;;;;;;;;;;;
 
@@ -158,8 +157,17 @@
 ;;;;;;;;;;;;;;;;;;;;;;;;;;;;; POSIT IMPLS ;;;;;;;;;;;;;;;;;;;;;;;;;;;
 
 (define-representation <posit8>  #:cost 1)
+(define-operation (if.p8 [c <bool>] [t <posit8>] [f <posit8>]) <posit8>
+  #:spec (if c t f) #:impl if-impl
+  #:cost 1 #:aggregate if-cost)
 (define-representation <posit16> #:cost 1)
+(define-operation (if.p16 [c <bool>] [t <posit16>] [f <posit16>]) <posit16>
+  #:spec (if c t f) #:impl if-impl
+  #:cost 1 #:aggregate if-cost)
 (define-representation <posit32> #:cost 1)
+(define-operation (if.p32 [c <bool>] [t <posit32>] [f <posit32>]) <posit32>
+  #:spec (if c t f) #:impl if-impl
+  #:cost 1 #:aggregate if-cost)
 
 (define-operations ([x <posit8>] [y <posit8>]) <bool>
   [==.p8 #:spec (== x y) #:impl posit8=  #:cost 1]

src/core/alt-table.rkt

@@ -37,9 +37,6 @@
         [(? symbol?) ((node-cost-proc node repr))]
         [(? number?) 0] ; specs
         [(approx _ impl) (vector-ref costs impl)]
-        [(list 'if cond ift iff)
-         (define cost-proc (node-cost-proc node repr))
-         (cost-proc (vector-ref costs cond) (vector-ref costs ift) (vector-ref costs iff))]
         [(list (? (negate impl-exists?) impl) args ...) 0] ; specs
         [(list impl args ...)
          (define cost-proc (node-cost-proc node repr))

src/core/bsearch.rkt

@@ -42,11 +42,13 @@
        (for/fold ([expr (alt-expr (list-ref alts (sp-cidx (last splitpoints))))])
                  ([splitpoint (cdr (reverse splitpoints))])
          (define repr (repr-of (sp-bexpr splitpoint) ctx))
+         (define if-impl (get-fpcore-impl 'if '() (list (get-representation 'bool) repr repr)))
          (define <=-impl (get-fpcore-impl '<= '() (list repr repr)))
-         `(if (,<=-impl ,(sp-bexpr splitpoint)
-                        ,(literal (repr->real (sp-point splitpoint) repr) (representation-name repr)))
-              ,(alt-expr (list-ref alts (sp-cidx splitpoint)))
-              ,expr)))
+         `(,if-impl (,<=-impl ,(sp-bexpr splitpoint)
+                              ,(literal (repr->real (sp-point splitpoint) repr)
+                                        (representation-name repr)))
+                    ,(alt-expr (list-ref alts (sp-cidx splitpoint)))
+                    ,expr)))
 
      ;; We don't want unused alts in our history!
      (define-values (alts* splitpoints*) (remove-unused-alts alts splitpoints))

src/core/compiler.rkt

@@ -46,9 +46,6 @@
     [(list op a b c) (op (vector-ref regs a) (vector-ref regs b) (vector-ref regs c))]
     [(list op args ...) (apply op (map (curry vector-ref regs) args))]))
 
-(define (if-proc c a b)
-  (if c a b))
-
 (define (batch-remove-approx batch)
   (batch-replace batch
                  (lambda (node)
@@ -79,7 +76,6 @@
                 ([node (in-vector (batch-nodes batch*) num-vars)])
       (match node
         [(literal value (app get-representation repr)) (list (const (real->repr value repr)))]
-        [(list 'if c t f) (list if-proc c t f)]
         [(list op args ...) (cons (impl-info op 'fl) args)])))
 
   (make-progs-interpreter vars instructions (batch-roots batch*)))

src/core/egg-herbie.rkt

@@ -294,10 +294,6 @@
       [`(Explanation ,body ...) `(Explanation ,@(map (lambda (e) (loop e type)) body))]
       [(list 'Rewrite=> rule expr) (list 'Rewrite=> (get-canon-rule-name rule rule) (loop expr type))]
       [(list 'Rewrite<= rule expr) (list 'Rewrite<= (get-canon-rule-name rule rule) (loop expr type))]
-      [(list 'if cond ift iff)
-       (if (representation? type)
-           (list 'if (loop cond (get-representation 'bool)) (loop ift type) (loop iff type))
-           (list 'if (loop cond 'bool) (loop ift type) (loop iff type)))]
       [(list (? impl-exists? impl) args ...) (cons impl (map loop args (impl-info impl 'itype)))]
       [(list op args ...)
        #:when (string-contains? (~a op) "unsound")
@@ -324,7 +320,7 @@
           (cons '(*.f64 x y) '(*.f64 $var0 $var1))
           (cons '(+.f64 (*.f64 x y) #s(literal 2 binary64)) '(+.f64 (*.f64 $var0 $var1) 2))
           (cons '(cos.f32 (PI.f32)) '(cos.f32 (PI.f32)))
-          (cons '(if (TRUE) x y) '(if (TRUE) $var0 $var1))))
+          (cons '(if.f64 (TRUE) x y) '(if.f64 (TRUE) $var0 $var1))))
 
   (let ([egg-graph (make-egraph-data)])
     (for ([(in expected-out) (in-dict test-exprs)])
@@ -540,9 +536,9 @@
     [(cons f _) ; application
      (cond
        [(eq? f '$approx) (platform-reprs (*active-platform*))]
-       [(eq? f 'if) (all-reprs/types)]
        [(string-contains? (~a f) "unsound") (list 'real)]
        [(impl-exists? f) (list (impl-info f 'otype))]
+       [(eq? f 'if) '(real bool)]
        [else (list (operator-info f 'otype))])]))
 
 ;; Rebuilds an e-node using typed e-classes
@@ -556,23 +552,15 @@
         (define spec (u32vector-ref ids 0))
         (define impl (u32vector-ref ids 1))
         (list '$approx (lookup spec (representation-type type)) (lookup impl type))]
-       [(eq? f 'if) ; if expression
-        (define cond (u32vector-ref ids 0))
-        (define ift (u32vector-ref ids 1))
-        (define iff (u32vector-ref ids 2))
-        (define cond-type
-          (if (representation? type)
-              (get-representation 'bool)
-              'bool))
-        (list 'if (lookup cond cond-type) (lookup ift type) (lookup iff type))]
        [(string-contains? (~a f) "unsound")
         (define op (string->symbol (string-replace (symbol->string f) "unsound-" "")))
         (list* op (map (λ (x) (lookup (u32vector-ref ids x) 'real)) (range (u32vector-length ids))))]
        [else
         (define itypes
-          (if (impl-exists? f)
-              (impl-info f 'itype)
-              (operator-info f 'itype)))
+          (cond
+            [(impl-exists? f) (impl-info f 'itype)]
+            [(eq? f 'if) (list 'bool type type)]
+            [else (operator-info f 'itype)]))
         ; unsafe since we don't check that |itypes| = |ids|
         ; optimize for common cases to avoid extra allocations
         (cons
@@ -840,10 +828,6 @@
             [(? number?) (platform-repr-cost (*active-platform*) type)]
             [(? symbol?) (platform-repr-cost (*active-platform*) type)]
             [(list '$approx x y) 0]
-            [(list 'if c x y)
-             (match (platform-if-cost (*active-platform*))
-               [`(max ,n) n] ; Not quite right
-               [`(sum ,n) n])]
             [(list op args ...) (impl-info op 'cost)])
           1))
     (values (string->symbol (format "~a.~a" n k))
@@ -1032,11 +1016,8 @@
                (representation-type type)
                type))
          (approx (loop spec spec-type) (loop impl type))]
-        [(list 'if cond ift iff)
-         (if (representation? type)
-             (list 'if (loop cond (get-representation 'bool)) (loop ift type) (loop iff type))
-             (list 'if (loop cond 'bool) (loop ift type) (loop iff type)))]
         [(list (? impl-exists? impl) args ...) (cons impl (map loop args (impl-info impl 'itype)))]
+        [(list 'if c t f) (list 'if (loop c 'bool) (loop t 'real) (loop f 'real))]
         [(list op args ...) (cons op (map loop args (operator-info op 'itype)))])))
 
   (define (eggref id)
@@ -1067,16 +1048,14 @@
              (define final-spec (egg-parsed->expr spec* spec-type))
              (define final-spec-idx (mutable-batch-munge! out final-spec))
              (approx final-spec-idx (loop impl type))]
-            [(list 'if (app eggref cond) (app eggref ift) (app eggref iff))
-             (if (representation? type)
-                 (list 'if (loop cond (get-representation 'bool)) (loop ift type) (loop iff type))
-                 (list 'if (loop cond 'bool) (loop ift type) (loop iff type)))]
             [(list (? impl-exists? impl) (app eggref args) ...)
              (define args*
                (for/list ([arg (in-list args)]
                           [type (in-list (impl-info impl 'itype))])
                  (loop arg type)))
              (cons impl args*)]
+            [(list 'if c t f)
+             (list 'if (loop (eggref c) 'bool) (loop (eggref t) type) (loop (eggref f) type))]
             [(list (? operator-exists? op) (app eggref args) ...)
              (define args*
                (for/list ([arg (in-list args)]
@@ -1120,7 +1099,6 @@
     [(? symbol?) 1]
     ; approx node
     [(list '$approx _ impl) (rec impl)]
-    [(list 'if cond ift iff) (+ 1 (rec cond) (rec ift) (rec iff))]
     [(list (? impl-exists? impl) args ...)
      (match (pow-impl-args impl args)
        [(cons _ e)
@@ -1150,9 +1128,6 @@
         ((node-cost-proc node repr))]
        ; approx node
        [(list '$approx _ impl) (rec impl)]
-       [(list 'if cond ift iff) ; if expression
-        (define cost-proc (node-cost-proc node type))
-        (cost-proc (rec cond) (rec ift) (rec iff))]
        [(list (? impl-exists?) args ...) ; impls
         (define cost-proc (node-cost-proc node type))
         (apply cost-proc (map rec args))])]

src/core/egglog-herbie-tests.rkt

@@ -51,10 +51,10 @@
 
   (check-equal? (e2->expr '(Var "y")) 'y)
 
-  (check-equal? (e2->expr '(IfTy (Var "y")
-                                 (Num (bigrat (from-string "1") (from-string "2")))
-                                 (Num (bigrat (from-string "0") (from-string "1")))))
-                '(if y 1/2 0))
+  (check-equal? (e2->expr '(Iff64Ty (Var "y")
+                                    (Num (bigrat (from-string "1") (from-string "2")))
+                                    (Num (bigrat (from-string "0") (from-string "1")))))
+                '(if.f64 y 1/2 0))
 
   (check-equal? (e2->expr '(Mulf64Ty (Num (bigrat (from-string "2") (from-string "1")))
                                      (Num (bigrat (from-string "3") (from-string "1")))))
@@ -154,10 +154,10 @@
                                                (Var "z"))))
                 '(*.f32 3/2 (+.f32 4/5 z)))
 
-  (check-equal? (e2->expr '(IfTy (Var "cond")
-                                 (Num (bigrat (from-string "7") (from-string "8")))
-                                 (Num (bigrat (from-string "-2") (from-string "3")))))
-                '(if cond 7/8 -2/3)))
+  (check-equal? (e2->expr '(Iff32Ty (Var "cond")
+                                    (Num (bigrat (from-string "7") (from-string "8")))
+                                    (Num (bigrat (from-string "-2") (from-string "3")))))
+                '(if.f32 cond 7/8 -2/3)))
 
 ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
 ;; Testing API
@@ -262,6 +262,8 @@
                                        (Neqf64Ty . !=.f64)
                                        (NotboolTy . not.bool)
                                        (OrboolTy . or.bool)
+                                       (Iff32Ty . if.f32)
+                                       (Iff64Ty . if.f64)
                                        (Pif32Ty . PI.f32)
                                        (Pif64Ty . PI.f64)
                                        (Powf32Ty . pow.f32)
@@ -311,7 +313,7 @@
            "../syntax/platform.rkt"
            "../utils/float.rkt"
            "../syntax/load-platform.rkt")
-  (activate-platform! (*platform-name*))
+  (activate-platform! "c")
 
   (define batch
     (progs->batch (list '(-.f64 (sin.f64 (+.f64 x eps)) (sin.f64 x))