Created
October 22, 2023 07:12
-
-
Save joinr/c001de6ca120689ad199f8ffa77a0ee6 to your computer and use it in GitHub Desktop.
primitive reducing, now primitive backed vectors don't have to suck as much.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
(ns prims) | |
(defprotocol IPrimitiveReducible | |
(reduce-prim- [this f init])) | |
(definterface IPrimitiveReduce | |
(^long LLLreduce [^clojure.lang.IFn$LLL f ^long init]) | |
(^double DDDreduce [^clojure.lang.IFn$DDD f ^double init])) | |
(def am->type | |
(let [m @#'clojure.core/ams] | |
(zipmap (vals m) (keys m)))) | |
(defn padd ^long [^long x ^long y] (unchecked-add x y)) | |
(defn long-red ^long [^longs xs ^clojure.lang.IFn$LLL f ^long init] | |
(let [l (alength xs)] | |
(loop [idx 0 | |
acc init] | |
(if (< idx l) | |
(recur (unchecked-inc-int idx) | |
(.invokePrim ^clojure.lang.IFn$LLL f acc (aget xs idx))) | |
acc)))) | |
(defn double-red ^double [^doubles xs ^clojure.lang.IFn$DDD f ^double init] | |
(let [l (alength xs)] | |
(loop [idx 0 | |
acc init] | |
(if (< idx l) | |
(recur (unchecked-inc-int idx) | |
(.invokePrim ^clojure.lang.IFn$DDD f acc (aget xs idx))) | |
acc)))) | |
;;^clojure.core.VecNode root | |
;;seems like we can walk the nodes until we hit a non-vec-node. | |
;;nil = empty. | |
(def longtype (Class/forName "[J")) | |
(def doubletype (Class/forName "[D")) | |
(defn long-array? [x] (instance? longtype x)) | |
(defn double-array? [x] (instance? doubletype x)) | |
(defn array-walk [^clojure.core.VecNode root leaf? f] | |
(when-let [^objects children (and root (.arr root))] | |
(if (leaf? children) | |
(f children) | |
(let [bound (alength children)] | |
(loop [idx 0] | |
(when (< idx bound) | |
(do (array-walk (aget children idx) leaf? f) | |
(recur (unchecked-inc idx))))))))) | |
(defn vec-array-walk [^clojure.core.Vec v leaf? f] | |
(do (array-walk (.root v) leaf? f) | |
(f (.tail v)))) | |
(defn vec-wrapper ^IPrimitiveReduce [^clojure.core.Vec coll] | |
(reify IPrimitiveReduce | |
(LLLreduce [this f init] | |
(let [^clojure.lang.IFn$OOLL | |
collapse long-red | |
res (long-array 1) | |
_ (aset res 0 init)] | |
(vec-array-walk coll long-array? (fn [xs] | |
(aset res 0 | |
(.invokePrim collapse xs f (aget res 0))))) | |
(aget res 0))) | |
(DDDreduce [this f init] | |
(let [^clojure.lang.IFn$OODD | |
collapse double-red | |
res (double-array 1) | |
_ (aset res 0 init)] | |
(vec-array-walk coll double-array? (fn [xs] | |
(aset res 0 | |
(.invokePrim collapse xs f (aget res 0))))) | |
(aget res 0))))) | |
(extend-protocol IPrimitiveReducible | |
Object | |
(reduce-prim- [this f init] | |
(reduce f init this)) | |
clojure.core.Vec | |
(reduce-prim- [this f init] | |
(let [wrapper (vec-wrapper this)] | |
(case (am->type (.am this)) | |
:long (.LLLreduce wrapper f (long init)) | |
(.DDDreduce wrapper f (double init)))))) | |
;;minor boxing on input... | |
(defn reduce-prim [f init coll] | |
(reduce-prim- coll f init)) | |
(comment ;;testing | |
(require '[criterium.core :as c]) | |
(def v (vec (range 100000))) | |
(def pv (into (vector-of :long) (range 100000))) | |
(c/quick-bench (reduce + 0 v)) | |
;; Execution time mean : 1.142201 ms | |
(c/quick-bench (reduce + 0 pv)) | |
;;Execution time mean : 3.476692 ms ouch, boxing! | |
(c/quick-bench (reduce-prim padd 0 pv)) | |
;;Execution time mean : 128.043662 ╡s 8.9x not bad... | |
) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment