extract + parsing to herbie IR

herbie-fp · Aug 28, 2024 · c2eca4e · c2eca4e
1 parent 1d7400b
commit c2eca4e
Showing 1 changed file with 59 additions and 58 deletions.
diff --git a/src/core/egg-herbie.rkt b/src/core/egg-herbie.rkt
@@ -319,52 +319,52 @@
       [(list (? impl-exists? impl) args ...) (cons impl (map loop args (impl-info impl 'itype)))]
       [(list op args ...) (cons op (map loop args (operator-info op 'itype)))])))
 
-#;(define (egg-batch->batch batch rename-dict type)
-    (define nodes-length (batch-nodes-length batch))
-    (define nodes (batch-nodes batch))
-    (define types (make-vector nodes-length type))
-
-    ; Rewriting function for node
-    (define (rewrite-node node type)
-      (match node
-        [(? number?) (if (representation? type) (literal node (representation-name type)) node)]
-        [(? symbol?)
-         ; To be checked with PI instruction
-         (if (hash-has-key? rename-dict node)
-             (car (hash-ref rename-dict node)) ; variable (extract uncanonical name)
-             node)] ; constant function
-        ; An error can occur here because spec is not a node, it is expr
-        [(list '$approx spec impl)
-         (define spec-type (if (representation? type) (representation-type type) type))
-         (vector-set! types impl type)
-         (approx (egg-parsed->expr spec rename-dict spec-type) impl)]
-        [(list 'if cond ift iff)
-         (match (representation? type)
-           [#t
-            (vector-set! types cond (get-representation 'bool))
-            (vector-set! types ift type)
-            (vector-set! types iff type)
-            (list 'if cond ift iff)]
-           [#f
-            (vector-set! types cond 'bool)
-            (vector-set! types ift type)
-            (vector-set! types iff type)
-            (list 'if cond ift iff)])]
-        [(list (? impl-exists? impl) args ...)
-         (for ([arg (in-list args)]
-               [type (in-list (impl-info impl 'itype))])
-           (vector-set! types arg type))
-         (cons impl args)]
-        [(list op args ...)
-         (for ([arg (in-list args)]
-               [type (in-list (operator-info op 'itype))])
-           (vector-set! types arg type))
-         (cons op args)]))
-
-    (for ([node (in-vector nodes (- nodes-length 1) -1 -1)] ; reversed over nodes
-          [type (in-vector types (- nodes-length 1) -1 -1)]
-          [n (in-range (- nodes-length 1) -1 -1)])
-      (vector-set! nodes n (rewrite-node node type))))
+(define (egg-batch->batch batch rename-dict type)
+  (define nodes-length (batch-length batch))
+  (define nodes (batch-nodes batch))
+  (define types (make-vector nodes-length type))
+
+  ; Rewriting function for node
+  (define (rewrite-node node type)
+    (match node
+      [(? number?) (if (representation? type) (literal node (representation-name type)) node)]
+      [(? symbol?)
+       ; To be checked with PI instruction
+       (if (hash-has-key? rename-dict node)
+           (car (hash-ref rename-dict node)) ; variable (extract uncanonical name)
+           node)] ; constant function
+      ; An error can occur here because spec is not a node, it is expr
+      [(list '$approx spec impl)
+       (define spec-type (if (representation? type) (representation-type type) type))
+       (vector-set! types impl type)
+       (approx (egg-parsed->expr spec rename-dict spec-type) impl)]
+      [(list 'if cond ift iff)
+       (match (representation? type)
+         [#t
+          (vector-set! types cond (get-representation 'bool))
+          (vector-set! types ift type)
+          (vector-set! types iff type)
+          (list 'if cond ift iff)]
+         [#f
+          (vector-set! types cond 'bool)
+          (vector-set! types ift type)
+          (vector-set! types iff type)
+          (list 'if cond ift iff)])]
+      [(list (? impl-exists? impl) args ...)
+       (for ([arg (in-list args)]
+             [type (in-list (impl-info impl 'itype))])
+         (vector-set! types arg type))
+       (cons impl args)]
+      [(list op args ...)
+       (for ([arg (in-list args)]
+             [type (in-list (operator-info op 'itype))])
+         (vector-set! types arg type))
+       (cons op args)]))
+
+  (for ([node (in-vector nodes (- nodes-length 1) -1 -1)] ; reversed over nodes
+        [type (in-vector types (- nodes-length 1) -1 -1)]
+        [n (in-range (- nodes-length 1) -1 -1)])
+    (vector-set! nodes n (rewrite-node node type))))
 
 ;; Parses a string from egg into a single S-expr.
 (define (egg-expr->expr egg-expr egraph-data type)
@@ -1169,6 +1169,7 @@
      (define batch (finalize-batch))
 
      ; -------------------------- Debooging
+     ; This debooging crashes if unparsed approx nodes exist in the batch
      #;(printf "roots#=~a, eclasses#=~a\n"
                (vector-length (batch-roots batch))
                (vector-length (vector-ref eclasses id*)))
@@ -1185,22 +1186,22 @@
 
      ; translate egg IR to Herbie IR
      (define egg->herbie (regraph-egg->herbie regraph))
-     #;(define exprs
-         (for/list ([egg-expr (in-list egg-exprs)])
-           (egg-parsed->expr (flatten-let egg-expr) egg->herbie type)))
+     (define exprs
+       (for/list ([egg-expr (in-list egg-exprs)])
+         (egg-parsed->expr (flatten-let egg-expr) egg->herbie type)))
 
      ; translate egg IR batch to Herbie IR batch
-     #;(egg-batch->batch batch egg->herbie type)
+     (egg-batch->batch batch egg->herbie type)
 
      ; -------------------------- Debooging
-     #;(define exprs* (batch->progs batch))
-     #;(when (not (equal? exprs* exprs))
-         (println (vector-ref eclasses id*))
-         (for ([expr* (in-list exprs*)]
-               [expr (in-list exprs)])
-           (printf "expr* = ~a\n" expr*)
-           (printf "expr  = ~a\n\n" expr)
-           (sleep 5)))
+     (define exprs* (batch->progs batch))
+     (when (not (equal? exprs* exprs))
+       (println (vector-ref eclasses id*))
+       (for ([expr* (in-list exprs*)]
+             [expr (in-list exprs)])
+         (printf "expr* = ~a\n" expr*)
+         (printf "expr  = ~a\n\n" expr)
+         (sleep 5)))
      ; -------------------------
 
      batch]