bsless / clj-fast

Unpredictably faster Clojure
Eclipse Public License 2.0
242 stars 1 forks source link

General inliner #29

Open bsless opened 3 years ago

bsless commented 3 years ago

All of my inlining implementations, although close to the original implementation they're mimicking, are ad-hoc. They require prior knowledge of the inlined form. Ideally, I would want to expose a macro like definline which will include a sophisticated :inline function. This inliner will examine the call site and try to inline what it can.

Prior works on this include F expressions https://github.com/halgari/heliotrope https://web.wpi.edu/Pubs/ETD/Available/etd-090110-124904/unrestricted/jshutt.pdf

And Haskell's Core compiler https://gitlab.haskell.org/ghc/ghc/-/wikis/commentary/compiler/core-to-core-pipeline

I started trying to generalize an implementation, starting out with replacing function application with a rewrite rule which would replace application by the instructions to write an application:

(defn abstract-call
  [sym]
  (fn [& args]
    `(~sym ~@args)))

(defmacro ac
  [sym]
  `(abstract-call ~sym))

This allows to quite trivially port an existing definition:

(defn a-get-in
  [m ks]
  (reduce (ac `get) m ks))

(a-get-in 'm '[a b c])
;=>
(clojure.core/get (clojure.core/get (clojure.core/get m a) b) c)

It requires some massaging of more complex definitions but essentially works:

(defn a-assoc-in
  [m [k & ks] v]
  (let [g (gensym)]
    `(let [~g ~m]
       ~(if ks
          ((ac `assoc) g k (a-assoc-in ((ac `get) g k) ks v))
          ((ac `assoc) g k v)))))

(a-assoc-in 'm '[a b c] 'v)
;=>
(clojure.core/let
 [G__11548 m]
 (clojure.core/assoc
  G__11548
  a
  (clojure.core/let
   [G__11549 (clojure.core/get G__11548 a)]
   (clojure.core/assoc
    G__11549
    b
    (clojure.core/let
     [G__11550 (clojure.core/get G__11549 b)]
     (clojure.core/assoc G__11550 c v))))))

Things get more complicated when trying to generalize this

This is my initial abortive attempt:

(defn replace-args
  [args form]
  (let [argm (zipmap args (repeatedly gensym))
        imap (interleave (vals argm) (keys argm))]
    `(let [~@imap]
       ~(walk/postwalk-replace argm form))))

(defn fnsym
  [sym]
  (when-let [v (resolve sym)]
    (when (and (ifn? (deref v)) (not (:macro (meta v))))
      (symbol v))))

(defn abstract-fn
  [sym]
  (if-let [sym (fnsym sym)]
    `(abstract-call ~sym)
    sym))

(comment
  ((abstract-fn 'get) 'm 'k))

(defn abstract-form
  [name form]
  (walk/postwalk
   (fn [expr]
     (if-let [expr (and (symbol? expr)
                        (not= name expr)
                        (abstract-fn expr))]
       expr
       expr))
   form))

(defn regenerate-form
  [name args form]
  (fn [& args']
    (let [gnosis (map (fn [argn arg] (if (known-at-callsite? arg)
                                       (with-meta argn (assoc (meta argn) :known true))
                                       argn)) args args')
          known (filter (comp :known meta) gnosis)
          unknown (remove (comp :known meta) gnosis)]
      (if (seq known)
        (->> form
             (replace-args unknown)
             (abstract-form name))
        'boring))))

(defn emit-inliner
  [name args form]
  (let [generator (regenerate-form name args form)]
    (fn [& callsite]
      (let [emitter (apply generator callsite)
            emitter (eval `(fn [~@args] ~emitter))]
        (apply emitter callsite)))))

(defmacro definline+
  [name & decl]
  (let [[pre-args [args expr]] (split-with (comp not vector?) decl)]
    `(do
       (defn ~name ~@pre-args ~args ~expr)
       (alter-meta! (var ~name) assoc :inline (emit-inliner (quote ~name) (quote ~args) (quote ~expr)))
       (var ~name))))

(definline+ my-assoc-in
  [m [k & ks] v]
  (if ks
    (assoc m k (my-assoc-in (get m k) ks v))
    (assoc m k v)))

(defn foo
  [m v]
  (my-assoc-in m [1 2 3] v))

There are heuristics to be considered when attempting to inline, some of which are discussed in the Haskell Core compiler literature, such as the number of occurrences of a variable in a form determining if it needs to be bound (see the input map in assoc-in).

cc: @joinr, what do you think, does a solution present itself readily here, or is it completely inscrutable?

joinr commented 3 years ago

Interesting. That would be a useful unification of the extant macro-based approach. I don't have anything to add at the moment, but will ponder the implications.

bsless commented 3 years ago

I think a very slow and inefficient way to do this can be a combination of beta reduction and constant propagation. Most of my inlining is just constant propagation for sequences until a fixed point is reached. A beta reduction is easy, just the following transformation:

((fn [x y z] ,,,) a b c)

(let [x a
      y b
      z c]
  ,,,)

Any thoughts on how I could perform / generalize constant propagation?

joinr commented 3 years ago

dredging up memories of constant folding from tim baldridge's presentation on tools.analyzer. Great presentation overall on learning the analyzer.

Beyond that, I don't have any grand ideas on compilation; some caveman square wheel ideas, but compilers are not my deepest area of knowledge.

bsless commented 3 years ago

You probably have better knowledge of them than I do. I'm wondering if I can perhaps sidestep the problem altogether and let something else do the term-rewriting and propagation for me. Thinking core.logic or meander.

joinr commented 3 years ago

Yeah, meander is probably custom made for this then. I am less familiar with it (although I did mess with some its evaluation engine once trying to optimize the emitted code a bit). core.logic is an interesting prospect as well; I've gotten in and out of it over the years (one time a couple of weeks ago).

bsless commented 3 years ago

I think tools.analyzer might be a better fit, it allows working with a AST and not raw Clojure forms. What I'd need is to figure out how to:

bsless commented 2 years ago

@joinr Took about a year of fumbling bug I think I hit on a workable approach https://github.com/bsless/clj-analyzer

joinr commented 2 years ago

That looks very promising. Any benchmarks for toy results?

bsless commented 2 years ago

Nothing yet, I only figured out the correct data model for occurrences' arithmetic yesterday. To get something meaningful I need two other things: case elimination and loop breakers.

joinr commented 2 years ago

Looks like term rewriting at the AST node level; using meander on top of analyzer. nice

bsless commented 2 years ago

I actually haven't brought meander in. Assuming inlining only in let bindings, this is all that you need to inline:

(defn do-inline
  [ast {:keys [name init]}]
  (ast/postwalk
   ast
   (fn [{:keys [op] :as ast}]
     (if (and (= op :local) (= name (:name ast)))
       init
       ast))))

(defn inline
  "Assume AST here is a let-form"
  [{:keys [bindings body] :as ast} binding]
  (let [bindings (remove #(identical? binding %) bindings)]
    (assoc
     ast
     :bindings (mapv #(do-inline % binding) bindings)
     :body (do-inline body binding))))

It turns out all the logic for testing whether I should inline wasn't the best fit for meander, and if you write this snippet in meander it's very verbose. I might use it for case elimination