Added in a fast version of game of life using compute tensors.

cnuernber · Jan 20, 2021 · 619315b · 619315b
1 parent 0c8017f
commit 619315b
Show file tree

Hide file tree

Showing 2 changed files with 58 additions and 6 deletions.
diff --git a/src/tech/v3/tensor/tensor_copy.clj b/src/tech/v3/tensor/tensor_copy.clj
@@ -90,8 +90,8 @@
                      (casting/host-flatten dst-dtype)
                      dst-dtype)
          ;;Getting the buffers defeats a check in the tensors
-         src-buffer (dtype-base/as-buffer (.buffer src))
-         dst-buffer (dtype-base/as-buffer (.buffer dst))]
+         src-buffer (when (.buffer src) (dtype-base/as-buffer (.buffer src)))
+         dst-buffer (when (.buffer dst) (dtype-base/as-buffer (.buffer dst)))]
      (when (and src-buffer
                 dst-buffer
                 (= src-dtype

diff --git a/test/tech/v3/apl/game_of_life.clj b/test/tech/v3/apl/game_of_life.clj
@@ -3,13 +3,15 @@
   (:require [tech.v3.tensor :as tens]
             [tech.v3.datatype :as dtype]
             [tech.v3.datatype.functional :as dtype-fn]
-            [clojure.test :refer :all]))
+            [tech.v3.datatype.bitmap :as bitmap]
+            [clojure.test :refer :all])
+  (:import [java.util Set]))
 
 
 (defn membership
   [lhs rhs]
-  (let [membership-set (set (dtype/->vector rhs))]
-    (dtype/emap #(contains? membership-set %) :boolean lhs)))
+  (let [membership-set (bitmap/->bitmap rhs)]
+    (dtype/emap #(.contains membership-set %) :boolean lhs)))
 
 
 (defn apl-take
@@ -50,7 +52,7 @@
     (tens/rotate tens offsets)))
 
 
-(def range-tens (tens/reshape (vec (range 9)) [3 3]))
+(def range-tens (tens/->tensor (partition 3 (range 9)) :datatype :int8))
 
 (def bool-tens (-> range-tens
                    (membership [1 2 3 4 7])
@@ -146,3 +148,53 @@
               (take 1000)
               last
               (tens/->jvm)))))
+
+
+(defmacro ^:private center-coord
+  [loc dshp]
+  `(let [loc# (rem ~loc ~dshp)]
+     (if (< loc# 0)
+       (+ loc# ~dshp)
+       loc#)))
+
+
+(defn game-of-life-compute-op
+  [input-tens]
+  (let [[yshp xshp] (dtype/shape input-tens)
+        yshp (long yshp)
+        xshp (long xshp)
+        input-tens (tens/ensure-tensor input-tens)]
+    (-> (tens/compute-tensor
+         [yshp xshp]
+         (fn [^long y ^long x]
+           (let [original (.ndReadLong input-tens y x)
+                 conv-sum (long (loop [idx 0
+                                       sum 0]
+                                  (if (< idx 9)
+                                    (let [conv-y (unchecked-dec (quot idx 3))
+                                          conv-x (unchecked-dec (rem idx 3))
+                                          y-coord (center-coord (+ y conv-y) yshp)
+                                          x-coord (center-coord (+ x conv-x) xshp)]
+                                      (recur (unchecked-inc idx)
+                                             (+ sum (.ndReadDouble input-tens y-coord
+                                                                   x-coord))))
+                                    sum)))]
+             (if (or (== 3 conv-sum)
+                     (and (not= 0 original)
+                          (== 4 conv-sum)))
+               1
+               0)))
+         (dtype/elemwise-datatype input-tens))
+        (dtype/clone))))
+
+(defn life-compute-seq
+  [input]
+  (cons input (lazy-seq (life-compute-seq (game-of-life-compute-op input)))))
+
+
+(deftest game-of-life-compute-test
+  (is (= end-state
+         (->> (life-compute-seq RR)
+              (take 1000)
+              last
+              (tens/->jvm)))))