Skip to content

Instantly share code, notes, and snippets.

@joinr
Created October 22, 2023 07:12
Show Gist options
  • Save joinr/c001de6ca120689ad199f8ffa77a0ee6 to your computer and use it in GitHub Desktop.
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.
(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