Neanderthal vs ND4J - vol 4 - Fast Vector Broadcasting in Java, CPU and CUDA

Even if I didn't know anything about linear algebra, I could employ a stock tool from the programmer's toolbox, iteration. I'll naively ask Neanderthal's matrix a to give me a sequence of its rows, which are vectors, and add x to those vectors one by one.

(defn naive-add-row-vector! [a x]
  (doseq [rw (rows a)]
    (axpy! x rw))
  a)

Let's test it briefly:

(let [a (fge 2 3 [1 2 3 4 5 6])
      x (entry! (fv 3) 1)]
  (naive-add-row-vector! a x))

#RealGEMatrix[float, mxn:2x3, layout:column, offset:0]
   ▥       ↓       ↓       ↓       ┓
   →       2.00    4.00    6.00
   →       3.00    5.00    7.00
   ┗                               ┛

As the phrase goes, you can "convince yourself that this is equivalent functionality" to Nd4j's. Neanderthal has somewhat prettier printing, but that is not of major concern right now.

Since we started with the most naive implementation, I expected Nd4j to take a giant lead here. I have been so naive! Neanderthal was not that much behind!

(defn bench-neanderthal-naive-add-row-vector [m n]
  (with-release [a (entry! (fge m n {:layout :row}) 1)
                 x (entry! (fv n) 1)]
    (quick-bench
     (do (naive-add-row-vector! a x)
         true))))

;;(bench-neanderthal-naive-add-row-vector 100 100)

Here are the results:

;; 10x15: 1.379596 µs
;; 100x100: 11.467549 µs
;; 1000x1000: 216.736544 µs
;; 100x10000: 140.238912 µs
;; 10000x100: 1.163332 ms
;; 10000x10000: 31.391642 ms

But why stop here? We can improve this naive implementation and make it less naive, to show you that with a good tool in your toolbox, and a little knowledge of textbook math you can make wonders.

So, what caught my eye in the previous measurements, is that cases when the matrix is "flat" require a lot of looping iterations, each calling native operation and wasting its time by not giving it enough work to do. What I would prefer is, of course, to call fewer more demanding operations, preferably just one :)

Take a look at the diagram on the Nd4j user guide that explains what broadcasting does, and recall a bit of basic linear algebra. If you imagine the vector to be broadcast to be a \(1 \times{} n\) matrix, and imagine a \(m \times{} 1\) matrix of ones, you may recognize that their product places the elements exactly at the places where they need to end up.

This might sound a bit messy. Whaaat? I'd have to convert the vector to a matrix and whatnot? Although it would be easy in Neanderthal, and it would be hidden inside a function, there is something even simpler. Yep, whatever I end up needing for these kinds of tasks, Neanderthal just end up already having a convenient function to use!

In this case, it is the rk! function, which can multiply each combination of elements of two orthogonal vectors resulting in a matrix that can be summed element-by-element with another matrix.

Here is a less naive broadcasting implementation in Neanderthal. This time, a two-liner:

(defn less-naive-add-row-vector! [a x]
  (with-release [y (entry! (raw (col a 0)) 1)]
    (rk! y x a)))

(let [a (fge 2 3 [1 2 3 4 5 6])
      x (entry! (fv 3) 1)]
  (less-naive-add-row-vector! a x))

#RealGEMatrix[float, mxn:2x3, layout:column, offset:0]
   ▥       ↓       ↓       ↓       ┓
   →       2.00    4.00    6.00
   →       3.00    5.00    7.00
   ┗                               ┛

It works correctly; what about the performance?

(defn bench-neanderthal-less-naive-add-row-vector [m n]
  (with-release [a (entry! (fge m n {:layout :row}) 1)
                 x (entry! (fv n) 1)]
    (quick-bench
     (do (less-naive-add-row-vector! a x)
         true))))

;; 10x15: 506.063330 ns
;; 100x100: 1.768599 µs
;; 1000x1000: 61.031548 µs
;; 100x10000: 155.047204 µs
;; 10000x100: 66.700431 µs
;; 10000x10000: 26.320592 ms

Now Neanderthal even took the lead!

With small-ish matrices it is 3-4 × faster, while with large and huge matrices it is on par with Nd4J. The only slip is with "flat" matrices, it is almost 3× slower here, and that's because it has to create a huge vector of 10000 ones for each invocation of the function. That's something that can be improved, but this is a makeshift implementation anyway. In "real" code, I use optimized broadcasting (sorry, I hasn't released that library as open source), and it works faster.

When we have large arrays and need speed, the real saver is GPU computing anyway. So why not try that instead?

(defn bench-nd4j-addiRowVector-cuda [^long m ^long n]
  (let [a (Nd4j/ones m n)
        x (Nd4j/ones n)]
    (quick-bench
     (do (.addiRowVector ^INDArray a ^INDArray x)
         (.sumNumber ^INDArray x)))))

;;(bench-nd4j-addiRowVector-cuda 10000 10000)

;; 10x15: 241.076577 µs
;; 100x100: 138.564805 µs
;; 1000x1000: 223.194717 µs
;; 100x10000: 289.878918 µs
;; 10000x100: 211.342593 µs
;; 10000x10000: 7.895556 ms

Except for the largest dimension, Nd4j's broadcasting runs not only slower than Nd4j's CPU implementation, not only slower that a less naive Neanderthal broadcasting on the CPU; it is slower than the most naive broadcasting implementation on the CPU.

That's not only Nd4j's fault, since broadcasting is simply not demanding enough to squeeze the juice out of the GPU. I mentioned that more as a warning to inexperienced readers to not expect the magic words GPU and CUDA to magically speed up the code.

Let's see Neanderthal.

(defn bench-neanderthal-less-naive-add-row-vector-cuda [m n]
  (with-release [a (entry! (cuge m n) 1)
                 x (entry! (cuv n) 1)]
    (quick-bench
     (do (less-naive-add-row-vector! a x)
         (asum x)))))

(init)
(current-context! (context (device 0)))
(set-engine!)
;;(bench-neanderthal-less-naive-add-row-vector-cuda 10000 10000)
(release (current-context))

;; 10x15: 52.111330 µs
;; 100x100: 52.281520 µs
;; 1000x1000: 81.112859 µs
;; 100x10000: 80.279229 µs
;; 10000x100: 74.671313 µs
;; 10000x10000: 2.601956 ms

As I mentioned, on top of Nd4j's summing synchronization handicap, it has its own array creating and populating with ones handicap. Interestingly, a less naive broadcasting implementation turned out to be 3 × faster than Nd4j's CUDA code! Most of that is in GPU communication.

Without asum, and with explicit ClojureCL synchronize!, all these 50s turn into 28 microseconds. That's exactly why those not-so-demanding operations are not (much) faster on the GPU: the actual computation is just a small part of the time. GPU's have huge throughput, but have much higher latency than the CPU code, even when there is not any data transfer.

The real puzzle here is where Nd4j wasted these 150 microseconds that it is slower than Neanderthal in most cases, or why it stashed 5 milliseconds compared to Neanderthal in that 100M matrix…