From aff9ac968840e9c86719fb613bd2ed3c39b9905c Mon Sep 17 00:00:00 2001
From: Andy Wingo <wingo@pobox.com>
Date: Thu, 26 Sep 2024 11:14:52 +0200
Subject: [PATCH] Run sigbits fixpoint based on use/def graph, not cfg

* module/language/cps/specialize-numbers.scm (sigbits-ref): New helper.
(invert-graph*): New helper.
(compute-significant-bits): When visiting a term changes computed
needed-bits for one of its definitions, we need to revisit the variables
that contributed to its result (the uses), because they might need more
bits as well.  Previously we were doing this by enqueueing predecessors
to the term, which worked if the uses were defined in predecessors, or
if all defining terms were already in the worklist, which is the case
without loops.  But with loops, when revisiting a term, you could see
that it causes sigbits to change, enqueue its predecessors, but then the
predecessors don't change anything and the fixpoint stops before
reaching the definitions of the variables we need.  So instead we
compute the use-def graph and enqueue defs directly.
---
 module/language/cps/specialize-numbers.scm | 120 ++++++++++-----------
 1 file changed, 54 insertions(+), 66 deletions(-)

diff --git a/module/language/cps/specialize-numbers.scm b/module/language/cps/specialize-numbers.scm
index b46919a40..4983326f6 100644
--- a/module/language/cps/specialize-numbers.scm
+++ b/module/language/cps/specialize-numbers.scm
@@ -286,6 +286,9 @@
       (and (type<=? type (logior &exact-integer &u64 &s64))
            (range->sigbits min max)))))
 
+(define (sigbits-ref sigbits var)
+  (intmap-ref sigbits var (lambda (_) 0)))
+
 (define significant-bits-handlers (make-hash-table))
 (define-syntax-rule (define-significant-bits-handler
                       ((primop label types out def ...) param arg ...)
@@ -297,24 +300,42 @@
 (define-significant-bits-handler ((logand label types out res) param a b)
   (let ((sigbits (sigbits-intersect3 (inferred-sigbits types label a)
                                      (inferred-sigbits types label b)
-                                     (intmap-ref out res (lambda (_) 0)))))
+                                     (sigbits-ref out res))))
     (intmap-add (intmap-add out a sigbits sigbits-union)
                 b sigbits sigbits-union)))
 (define-significant-bits-handler ((logand/immediate label types out res) param a)
   (let ((sigbits (sigbits-intersect3 (inferred-sigbits types label a)
                                      param
-                                     (intmap-ref out res (lambda (_) 0)))))
+                                     (sigbits-ref out res))))
     (intmap-add out a sigbits sigbits-union)))
 
 (define (significant-bits-handler primop)
   (hashq-ref significant-bits-handlers primop))
 
+(define (invert-graph* defs)
+  "Given a graph LABEL->VAR..., return a graph VAR->LABEL....  Like the one
+in (language cps graphs), but different because it doesn't assume that
+the domain will be the same before and after."
+  (persistent-intmap
+   (intmap-fold (lambda (label vars out)
+                  (intset-fold
+                   (lambda (var out)
+                     (intmap-add! out var (intset label) intset-union))
+                   vars
+                   out))
+                defs
+                empty-intmap)))
+
 (define (compute-significant-bits cps types kfun)
   "Given the locally inferred types @var{types}, compute a map of VAR ->
 BITS indicating the significant bits needed for a variable.  BITS may be
 #f to indicate all bits, or a non-negative integer indicating a bitmask."
-  (let ((preds (invert-graph (compute-successors cps kfun))))
-    (let lp ((worklist (intmap-keys preds)) (out empty-intmap))
+  (let ((cps (intmap-select cps (compute-function-body cps kfun))))
+    ;; Label -> Var...
+    (define-values (defs uses) (compute-defs-and-uses cps))
+    ;; Var -> Label...
+    (define defs-by-var (invert-graph* defs))
+    (let lp ((worklist (intmap-keys cps)) (out empty-intmap))
       (match (intset-prev worklist)
         (#f out)
         (label
@@ -322,69 +343,36 @@ BITS indicating the significant bits needed for a variable.  BITS may be
            (define (continue out*)
              (if (eq? out out*)
                  (lp worklist out)
-                 (lp (intset-union worklist (intmap-ref preds label))
+                 (lp (intset-fold
+                      (lambda (use worklist)
+                        (intset-union worklist (intmap-ref defs-by-var use)))
+                      (intmap-ref uses label)
+                      worklist)
                      out*)))
-           (define (add-def out var)
-             (intmap-add out var 0 sigbits-union))
-           (define (add-defs out vars)
-             (match vars
-               (() out)
-               ((var . vars) (add-defs (add-def out var) vars))))
-           (define (add-unknown-use out var)
+           (define (add-unknown-use var out)
              (intmap-add out var (inferred-sigbits types label var)
                          sigbits-union))
-           (define (add-unknown-uses out vars)
-             (match vars
-               (() out)
-               ((var . vars)
-                (add-unknown-uses (add-unknown-use out var) vars))))
+           (define (default)
+             (intset-fold add-unknown-use (intmap-ref uses label) out))
            (continue
             (match (intmap-ref cps label)
-              (($ $kfun src meta self)
-               (if self (add-def out self) out))
-              (($ $kargs names vars term)
-               (let ((out (add-defs out vars)))
-                 (match term
-                   (($ $continue k src exp)
-                    (match exp
-                      ((or ($ $const) ($ $prim) ($ $fun) ($ $const-fun)
-                           ($ $code) ($ $rec))
-                       ;; No uses, so no info added to sigbits.
-                       out)
-                      (($ $values args)
-                       (match (intmap-ref cps k)
-                         (($ $kargs _ vars)
-                          (fold (lambda (arg var out)
-                                  (intmap-add out arg (intmap-ref out var (lambda (_) 0))
-                                              sigbits-union))
-                                out args vars))
-                         (($ $ktail)
-                          (add-unknown-uses out args))))
-                      (($ $call proc args)
-                       (add-unknown-use (add-unknown-uses out args) proc))
-                      (($ $callk label proc args)
-                       (let ((out (add-unknown-uses out args)))
-                         (if proc
-                             (add-unknown-use out proc)
-                             out)))
-                      (($ $calli args callee)
-                       (add-unknown-uses (add-unknown-use out callee) args))
-                      (($ $primcall name param args)
-                       (let ((h (significant-bits-handler name)))
-                         (if h
-                             (match (intmap-ref cps k)
-                               (($ $kargs _ defs)
-                                (h label types out param args defs)))
-                             (add-unknown-uses out args))))))
-                   (($ $branch kf kt src op param args)
-                    (add-unknown-uses out args))
-                   (($ $switch kf kt src arg)
-                    (add-unknown-use out arg))
-                   (($ $prompt k kh src escape? tag)
-                    (add-unknown-use out tag))
-                   (($ $throw src op param args)
-                    (add-unknown-uses out args)))))
-              (_ out)))))))))
+              (($ $kargs _ _ ($ $continue k _ ($ $primcall op param args)))
+               (match (significant-bits-handler op)
+                 (#f (default))
+                 (h
+                  (match (intmap-ref cps k)
+                    (($ $kargs _ defs)
+                     (h label types out param args defs))))))
+              (($ $kargs _ _ ($ $continue k _ ($ $values args)))
+               (match (intmap-ref cps k)
+                 (($ $kargs _ vars)
+                  (fold (lambda (arg var out)
+                          (intmap-add out arg (sigbits-ref out var)
+                                      sigbits-union))
+                        out args vars))
+                 (($ $ktail)
+                  (default))))
+              (_ (default))))))))))
 
 (define (specialize-operations cps)
   (define (u6-parameter? param)
@@ -416,7 +404,7 @@ BITS indicating the significant bits needed for a variable.  BITS may be
     (define (all-u64-bits-set? var)
       (operand-in-range? var &exact-integer (1- (ash 1 64)) (1- (ash 1 64))))
     (define (only-fixnum-bits-used? var)
-      (let ((bits (intmap-ref sigbits var)))
+      (let ((bits (sigbits-ref sigbits var)))
         (and bits (= bits (logand bits (target-most-positive-fixnum))))))
     (define (fixnum-result? result)
       (or (only-fixnum-bits-used? result)
@@ -429,7 +417,7 @@ BITS indicating the significant bits needed for a variable.  BITS may be
                        min max
                        (target-most-positive-fixnum)))))))
     (define (only-u64-bits-used? var)
-      (let ((bits (intmap-ref sigbits var)))
+      (let ((bits (sigbits-ref sigbits var)))
         (and bits (= bits (logand bits (1- (ash 1 64)))))))
     (define (u64-result? result)
       (or (only-u64-bits-used? result)
@@ -490,7 +478,7 @@ BITS indicating the significant bits needed for a variable.  BITS may be
             u64->fixnum)
            ((only-fixnum-bits-used? result)
             (lambda (cps k src u64)
-              (u64->fixnum/truncate cps k src u64 (intmap-ref sigbits result))))
+              (u64->fixnum/truncate cps k src u64 (sigbits-ref sigbits result))))
            (else
             u64->scm)))))
     (define (box-f64 result)
@@ -576,7 +564,7 @@ BITS indicating the significant bits needed for a variable.  BITS may be
            (('logand/immediate (? u64-result?) param a)
             (specialize-unop cps k src 'ulogand/immediate
                              (logand param
-                                     (or (intmap-ref sigbits result) -1)
+                                     (or (sigbits-ref sigbits a) -1)
                                      (1- (ash 1 64)))
                              a
                              (unbox-u64/truncate a) (box-u64 result)))
-- 
2.34.1

