Skip to content

Commit

Permalink
fmt
Browse files Browse the repository at this point in the history
  • Loading branch information
@pavpanchekha
pavpanchekha committed Aug 30, 2024
1 parent 360847b commit b590f9e
Show file tree
Hide file tree
Showing 2 changed files with 48 additions and 55 deletions.
81 changes: 41 additions & 40 deletions egg-herbie/main.rkt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -241,11 +241,10 @@
->
[f : _rust/string]
->
(values
(if (zero? (u32vector-length v))
(or (string->number f) (string->symbol f))
(string->symbol f))
v)))
(values (if (zero? (u32vector-length v))
(or (string->number f) (string->symbol f))
(string->symbol f))
v)))

;; This fairly long and ugly method is the core of our egg FFI.
;; It iterates over all enodes in the egraph, returning for each:
Expand Down Expand Up @@ -278,8 +277,7 @@
(define num-eclasses (u32vector-length eclass-ids))

(cond
[(= num-eclasses 0)
(in-range 0)]
[(= num-eclasses 0) (in-range 0)]
[else
(define first-eclass (u32vector-ref eclass-ids 0))
(define first-eclass-size (egraph_eclass_size egg-ptr first-eclass))
Expand All @@ -291,39 +289,42 @@

(make-do-sequence
(lambda ()
(values
(lambda (i)
(match-define (vector eclass-idx eclass-id eclass-size enode-idx) i)
(define node-size (egraph_enode_size egg-ptr eclass-id enode-idx))
(define vec
(match node-size [0 0-vec] [1 1-vec] [2 2-vec] [3 3-vec] [n (make-u32vector n)]))
(define-values (op args) (egraph_get_node egg-ptr eclass-id enode-idx vec))
(for ([i (in-range (u32vector-length args))])
(define sub-eclass-id (u32vector-ref args i))
(u32vector-set! args i (u32vector-ref egg-id->idx sub-eclass-id)))
(values eclass-idx op args))
(lambda (i)
(match-define (vector eclass-idx eclass-id eclass-size enode-idx) i)
(define enode-idx* (add1 enode-idx))
(cond
[(< enode-idx* eclass-size)
(vector-set! i 3 enode-idx*)]
[else
(define eclass-idx* (add1 eclass-idx))
(vector-set! i 0 eclass-idx*)
(when (< eclass-idx* num-eclasses)
(define eclass-id* (u32vector-ref eclass-ids eclass-idx*))
(vector-set! i 1 eclass-id*)
(define eclass-size* (egraph_eclass_size egg-ptr eclass-id*))
(vector-set! i 2 eclass-size*)
(vector-set! i 3 0))])
i)
(vector 0 first-eclass first-eclass-size 0)
(lambda (i)
(match-define (vector eclass-idx eclass-id eclass-size enode-idx) i)
(< eclass-idx num-eclasses))
#f
#f)))]))
(values (lambda (i)
(match-define (vector eclass-idx eclass-id eclass-size enode-idx) i)
(define node-size (egraph_enode_size egg-ptr eclass-id enode-idx))
(define vec
(match node-size
[0 0-vec]
[1 1-vec]
[2 2-vec]
[3 3-vec]
[n (make-u32vector n)]))
(define-values (op args) (egraph_get_node egg-ptr eclass-id enode-idx vec))
(for ([i (in-range (u32vector-length args))])
(define sub-eclass-id (u32vector-ref args i))
(u32vector-set! args i (u32vector-ref egg-id->idx sub-eclass-id)))
(values eclass-idx op args))
(lambda (i)
(match-define (vector eclass-idx eclass-id eclass-size enode-idx) i)
(define enode-idx* (add1 enode-idx))
(cond
[(< enode-idx* eclass-size) (vector-set! i 3 enode-idx*)]
[else
(define eclass-idx* (add1 eclass-idx))
(vector-set! i 0 eclass-idx*)
(when (< eclass-idx* num-eclasses)
(define eclass-id* (u32vector-ref eclass-ids eclass-idx*))
(vector-set! i 1 eclass-id*)
(define eclass-size* (egraph_eclass_size egg-ptr eclass-id*))
(vector-set! i 2 eclass-size*)
(vector-set! i 3 0))])
i)
(vector 0 first-eclass first-eclass-size 0)
(lambda (i)
(match-define (vector eclass-idx eclass-id eclass-size enode-idx) i)
(< eclass-idx num-eclasses))
#f
#f)))]))

;; egraph -> id -> id
(define-eggmath egraph_find (_fun _egraph-pointer _uint -> _uint))
Expand Down
22 changes: 7 additions & 15 deletions src/core/egg-herbie.rkt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -559,10 +559,10 @@
(define itypes (if (impl-exists? f) (impl-info f 'itype) (operator-info f 'itype)))
; unsafe since we don't check that |itypes| = |ids|
; optimize for common cases to avoid extra allocations
(cons
f
(for/list ([id (in-u32vector ids)] [itype (in-list itypes)])
(lookup id itype)))]))
(cons f
(for/list ([id (in-u32vector ids)]
[itype (in-list itypes)])
(lookup id itype)))]))

;; Rebuilds an e-node using typed e-classes
(define (rebuild-enode enode type lookup)
Expand Down Expand Up @@ -610,7 +610,8 @@

(define types (all-reprs/types))
(define type->idx
(for/hasheq ([type (in-list types)] [idx (in-naturals)])
(for/hasheq ([type (in-list types)]
[idx (in-naturals)])
(values type idx)))
(define num-types (hash-count type->idx))

Expand All @@ -632,24 +633,15 @@
; NOTE: nodes in typed eclasses are reversed relative
; to their position in untyped eclasses
(for ([(eclass op args) (in-egraph-enodes egg-ptr)])
#;(define enode
(if (or (number? op)
(and (hash-has-key? egg->herbie op) (zero? (u32vector-length args))))
op
(cons op args)))
; get all possible types for the enode
; lookup its correct eclass and add the rebuilt node
(define types (enode-type op egg->herbie))
(for ([type (in-list types)])
(define id (idx+type->id eclass type))
#;(define enode** (rebuild-enode enode type idx+type->id))
(define enode* (rebuild-enode* op args type egg->herbie idx+type->id))
#;(unless (equal? enode* enode**)
(error 'rebuild-enode "Don't match" enode* enode**))
(vector-set! id->eclass id (cons enode* (vector-ref id->eclass id)))
(match enode*
[(list _)
(vector-set! id->leaf? id #t)]
[(list _) (vector-set! id->leaf? id #t)]
[(list _ ids ...)
(for ([child-id (in-list ids)])
(vector-set! id->parents child-id (cons id (vector-ref id->parents child-id))))]
Expand Down

0 comments on commit b590f9e

Please sign in to comment.