Optimized Fractions Skill Score computation using numba

[nikeethr]

Parent issue: #188

Currently the numba implementation (although mostly complete) is lumped together in an example in #266

This issue is to separate out numba optimization of the fss score into it's own branch/PR, for easier review & tracking.

[nikeethr]

I might wait on #266 to be merged in before starting a proper PR for this since I expect merge conflicts otherwise. In the meantime I’ll work on an intermediate branch before cherry picking changes into this issue’s PR.

[nikeethr]

The parallel numba implementation is currently slower than the fast-fss implementation (i.e. pure numpy) during the first (pitch-forks down please) computation, after which it gets a speed up of 15-20x on subsequent executions. It is important to note that, numpy does already release the GIL for several basic functions. The advantage of using numba is basically to make this more explicit within python itself, without having to define optimized ufuncs in C along with some other cool benefits that I probably don't know about.

The main thing slowing the first kernel execution down is probably the parallelization of the FSS score. This part is not trivial and highly depends on the window size/position. Let's look at the following example:

---

                                                                                             
                                                                                             
  Requires thread synchronization          This is fine, all the data is in thread 2         
                ▲                                           ▲                                
                │                                           │                                
                │                                           │                                
                │                                           │                                
      ┌─────┬───┼─┬─────┬─────┐                   ┌─────┬───┼─┬─────┬─────┐                  
      │     │   │ │     │     ├─────► thread 1    │     │   │ │     │     ├─────► thread 1   
      │     │  A++│+B   │     │                   │     │   │ │     │     │                  
      ├─────┼───++│+────┼─────┤                   ├─────┼───┴─┴─────┼─────┤                  
      │     │  C++│+D   │     ├─────► thread 2    │     │  A+++++B  │     ├─────► thread 2   
      │     │     │     │     │                   │     │  C+++++D  │     │                  
      ├─────┼─────┼─────┼─────┤                   ├─────┼─────┬─────┼─────┤                  
      │     │     │     │     ├─────► thread 3    │     │     │     │     ├─────► thread 3   
      │     │     │     │     │                   │     │     │     │     │                  
      └─────┴─────┴─────┴─────┘                   └─────┴─────┴─────┴─────┘                  
                                                                                                                                                            

Let's say we only parallelize on rows for simplicity. So now each row potentially gets a thread or compute unit. Let's also put it out there that sharing data/memory/synchronizing between threads is expensive.

Its easy to see if the windows fit in nice non-overlapping 2-D grids, then this is a trivial parallelization, because the data is easily localized. Unfortunately, our windows have stride 1, which means that they are highly overlapping.

In the example above, non-overlapping meshgrid is shown as a reference point (i.e. stride length = window size).

On the left scenario in the ascii diagram, we're computing an arbitrary window (+) from a summed area table. We note that the window spans two threads, which means we need to share memory; because while the corners "A" and "B" live on thread 1, "C" and "D" live on thread 2. Yes technically we only need to synchronize 4 points... multiplied by the number of overlapping windows... and suddenly its not so technical...

On the second (right) scenario, all the computation happens in the same thread. So it doesn't matter that the window (+) overlaps two non-overlapping adjacent grids, because the thread has access to all the data - this we can get along with.

---

So what?

So back to numba, the current numba implementation parallelises everything pretty well (such as the summed area table) except for this last fss computation. We (I) just assume numba will "work it out". The reality is this is not an easy problem to optimize. As you may already be able to see, it depends on various factors such as. window size, number of available threads, memory synchronization strategy, data size. In fact it may not be worth doing any parallel compute at all!

To address this we have some options:

1. the sequential implementation is still quite fast, but we may not gain a 15-20x advantage over the numpy implementation, it'll be more around the 2-3x mark. Since the final FSS computation will always be the bottleneck, no matter how fast the summed area table computation is.
2. We could use a different algorithm.

---

Different algorithm?

Well.. it's actually not that different...

Note that the summed-area table only has an advantage if the same fields are looked up several times. In this case the window corners are looked up (4 points) over the entire field (N x M times). In most scenarios, we only need to perform the entire set of window look ups once per non-spatial dimension, as the data will change (e.g. we can't use the same summed-area table for T=0 and T=1, unless we do a summed-volume table, but what if we also include lead-time? Now the problem get's non-trivial...).

To put it plainly, the numba implementation doesn't have to be tied to the idea of using a summed area table. In fact, it may actually be slowing it down.

Here's an alternative strategy: What if like in the diagram above we decompose the 2D array into 1 thread per row. Now we can do a "vectorized" window operation along the rows quite easily, and it's trivially parallel. Since our window size is 1xW (W = width of window) there can be no vertical overlap, so the first (left) scenario does not exist. The accumulated result will give us a Nx1 vector (N = number of rows).

Then we repeat that over the columns, this time one thread per colu... but wait we only have 1 column, so actually this part probably doesn't need to be parallelized (even though there are strategies for this too, we've done the hard part by reducing 1 dimension in parallel - so lets not make life harder).

Now, the final accumulation just involves a 1-D sliding operation (window size = Hx1) along the Nx1 rows. If we follow the strategy of keeping a "decomposed fss score" as per #188, we now have an accumulated tuple of 3 numbers, the summed obs power, summed fcst power and the summed difference power, which can then be used to recompose the final scalar fss value.

Do I know this will work? Nope. But intuitively summation operations are associative, so in theory it should work.

---

Proof

TBD.

...or we just be lazy, try it out and verify by simulation!

[nikeethr]

Correction to above hypothesis

So there is a middle step I missed in the above idea for a different algorithm.

basically when computing the partial sums at each row - we can’t reduce it to 1 dimension (yet). We still need to keep track of the sliding sums. The reason is that the power operation itself is not associative.

Once we have the sliding sums along the rows, we should still have a NxM array (assuming appropriate prior zero padding). But this time each element i,j represents the sum from i,j -> i,j+W in the original field. W is the window width.

We can then repeat the same over the columns, this time performing the ^2 operation at the end. Again, we should still have a NxM but this time each element i,j should (in theory) represent the squared sum from i,j -> i+H,j+W in the original field, which is the squared sum for a particular window.

Now we can use the algorithm described earlier to reduce this to a single value - by computing it for each decomposed fss component before the final calculation.

---

How is this different to the summed area table?

The summed area table does not take into account the window size; it just computes the sum area table for an effective 1x1 window and then the resulting table is used for the window lookup - hence the thread synchronization issue.

This method, on the other hand, lumps the window lookup and power computation within a row/column making all the computations localised.

Again, this needs to be tested/proved.

---

One caveat however, is that the 1-D sliding windowed sums are about 2x slower than the integral image method (because, we need to perform 2 operations instead, one to add the suffix, and one to remove the prefix as the window traverses; in the summed area table we only need to add the previous element to the current element). However, the final FSS score will be computed at an order of magnitude faster, depending on the number of threads.

For small data sizes this may not provide much speed up.

Having said that, this is the perfect problem for a GPU because the data can be easily localised, compared to the previous algorithm. And should in theory avoid the kernel initialization hiccup, because there shouldn't be much thread synchronization going on, if any.

(i.e. if designed correctly we should be able to slap a cuda decorator or similar)

[nikeethr]

Semi-Proof

Here is a semi proof in order to illustrate the points to lessen the room for interpretation. We also assume most statements regarding algorithmic performance is for parallel compute pertaining to the current "Integral Image" and the proposed "Decomposed 1-D Sliding Sum" method only, and may not hold true for sequential cases.

Let us represent the fractions skills score as a composition of the average (or summed) square of:

• forecast across all sliding windows.
• observations across all sliding windows.
• the difference between obs and fcst across all sliding windows.

Since the components are independently summed, we only need to tackle one and generalize it to the other two. So, without loss of generality let's define the square sum of forecasts across all sliding windows as:

Eqn 0: FSS component: forecast

$$f_{sum} = {\sum_{i=0}^{N-1} \sum_{j=0}^{M-1} \left[ \sum_{k=0}^{H-1} \sum_{l=0}^{W-1} f[i+k,j+l] \right]^{2}}$$

where,
f[i,j] = forecast at grid point (i, j)
N = height of forecast field
M = width of forecast field
H = window height
W = window width
f_sum = accumulated sum of squares over all windows 

---

Because we can't move the outer sums into the "square" (associativity does not hold), we have to deal with the inner sums first.

The following function is defined to represent the result of running a 1-D sliding window sum along each row, let's also assume for the sake of simplicity that the domain is 0-padded, so that the equations are a bit simplified:

Eqn 1: Sums of sliding windows across rows

$$f_r[i,j] = \sum_{l=0}^{W-1} f[i,j+l]$$

We can see that the above sum is only computed over values of $l$, which only depends on (is relative to) $j$. In other words, we can calculate the sum for each $i$ (row) separately or parallelly.

We can do the same for columns. However, now suppose instead of summing over the columns of the original forecast field, we pass in $f_r$ above.

Eqn 2: Sums of sliding windows across columns

$$f_c[i,j] = \sum_{k=0}^{H-1} f_r[i+k,j]$$

Like Eqn 1, here $k$ is only dependent on $i$, and hence can be parallelly computed for each $j$. If we expand the Eqn 2 and square it we get:

Eqn 3: Squared sum of sliding windows across rows and columns

$$f_c^2[i,j] = \left[\sum_{k=0}^{H-1}\sum_{l=0}^{W-1} f[i+k,j+l]\right]^{2}$$

We see that the composed square sum $(f_c . f_r)^2$ is basically the inner sum in Eqn 0.

The outer sum is then trivial because we now have a N x M field of values given by $f_c^2[i,j]$, where each grid point is independent. In contrast, in the "Integral Image" algorithm in fast-fss we would have a Summed-area table by now and would have to calculate the squares by using a lookup method for the 4 window corners - which is highly likely to cause thread synchronization overheads as mentioned earlier.

This concludes our proof that decomposing the sliding window sums should give us the same result as the integral image method (which is not gone through here but is available in the reference paper)

---

Corollary 1: FSS computation using composed 1-D sliding windows is faster than Integral image

The number of operations for the proposed method here is O(N x M), albeit there are at least 4 of these computations per decomposed FSS value. However, since all the operations are trivially parallel as shown earlier, this can be reduced to O(N x M / P) where P is the number of compute units.

On the other hand, for the integral image, we would have to compute 2 partial sums for each decomposed FSS value, followed by the final FSS computation. Which in general is O(N_p x M_p / P + (N - N_p) x (M - M_p) + S) where, N_p x M_p represents the number of windows that can be trivially parallelized without any synchronization, N - N_p and M - M_p are the number of windows that overlap multiple threads and are generally not trivially parallelizable. S is an arbitrary synchronization overhead, which is highly dependent on various factors.

It's easy to show that the integral image method reaches the same gains as the sliding sum method only when the window size is 1x1. Furthermore, the single threaded case O(M x N) is not necessarily the worst case scenario as it may even perform better than the parallel method in some scenarios.

If we assume single threaded integral image operations are the average case, then we have a factor of P speedup with the proposed algorithm.

---

Corollary 2: The computation of 1-D sliding windows requires 2x the operations as Summed-area table

As mentioned in an earlier comment, the bottleneck in the FSS computation is the final FSS step that looksup the 4 corners. While, the 1-D sliding window method improves on it for the final score, it does come at a trade off compared to the Summed-area table method.

To see this , for any given element, in order to compute the next sum we just take the sum of the previous accumulated sum and the current element i.e.

$$f[i] = x[i] + f[i-1]$$

Whereas to compute the next sliding-window sum we'll also need to subtract the prefix sum from the window

$$f[i] = x[i] + f[i-1] - x[i-L]$$

where $L$ is the length of the window.

We see here that we have 1 binary operation and 2 lookups in the summed-area table, but 2 binary operations and 3 lookups in the sliding-sum operation. However, this is not an order of magnitude difference, and hence for larger datasets, is a smaller factor compared to the difference in compute of the FSS itself.

---

Finally some orders of operations. We can summarize the various algorithms as follows (for a N x N dataset and W x W window size):

• Naive 2-D sliding window: O(N^2 x W^2): very slow unless window size = 1x1, which is redundant.
• FFT: O(N^2 log(N))
• Summed-area table/Integral image: O(k1 x N^2), where k1 is a constant number of caching operations
• Decomposed sliding 1-D sum: O(k2 x N^2 / P), where k2 is a constant number of caching operations and P is the number of threads

In general k2 > k1, but in HPC environments or GPUs P is much, much larger

---

All of the above still needs to be verified via simulations in case of any logical missteps...

In summary,

• for cpu with limited threads: use the existing "Integral Image" method, except for the last window lookup step which should be done on a single thread for predictable performance.
• for gpu/hpc implementation. Use the "1-D Decomposed Sliding Window" method, outlined here.

Note: all of the above only apply to 2-D fields where we can decompose column and row operations, to maximise parallelisation. This will not work as well with 1-D fields, but there are other vectorized methods (by effectively buffering/chunking the data into a 2D field) that can apply to 1D that are beyond scope.

[nikeethr]

Generally, we want to make sure that the numba implementations can support GPUs where possible, without major future redesign. Especially since any usage of numba almost always implies that there is some optimization requirement of the underpinned algorithm.

Compute devices in general have the following components:

• global memory
• local memory
• global work size
• local work size

While global and local memory are hard limits, work sizes are design choices usually some multiple of the basic unit e.g. 32 or 64 and are used to divide global/local buffers into work items (which execute in parallel). Local memory (and size) is generally small, but quickly accessible between work items in a "work group" or compute unit (core). Global memory is large but expensive to access.

Computations are trivial when the 1-D sliding sum has a window size that fits inside local memory and its associated constraints. But we cannot always guarantee this.

---

Corollary 3: The 1-D sliding window computations can be inferred from the 2-D summed area table.

I've shown in #270 by implementation that the opencl versions of the 2-D summed area table are trivially parallelizable for most devices. So we'll take this as an assumption.

The summed area table for the forecast field (again without lost of generality) can be written as:

$$f_{sat}[i,j] = \sum_{k=0}^{i}\sum_{l=0}^{j} f[k,l]$$

Now suppose we define:

$$a_{ij} = f_{sat}[i,j]$$ $$b_{ij} = f_{sat}[i+H,j]$$ $$c_{ij} = f_{sat}[i,j+W]$$ $$d_{ij} = f_{sat}[i+H,j+W]$$

Without going into details, it can be shown that the squared window sum can be derived as

$$d_{ij} - c_{ij} - b_{ij} + a_{ij} = \sum_{k=0}^{H-1}\sum_{l=0}^{W-1} f[i+k,j+l]$$

Which is the same as $f_c[i,j]$ in Eqn 3.

---

In order to use Corollary 3 in practice, we can utilize a few things:

• Single instruction multiple data (SIMD): We can use vector types like int4 to store our data to take full advantage of the work sizes. Most GPUs come with wide SIMDs (128/256 bit).
• local buffer 1: int4, with, s0 = a, s1 = b, s2 = c, s3 = d
• local buffer 2: int4, with, s0 = (d_0 - c_0 - b_0 + a_0)^2, s1 = (d_1 - c_1 - b_1 + a_1)^2, ...

Here is how it would work:

• move data from global buffer to local buffers in the format described in local buffer (one time)
• Noting that local buffer 1 and local buffer 2 shared in each compute unit (core):
  1. compute the square sums of 4 adjacent entries of local buffer 1 and store them in the SIMD of local buffer 2
  2. compute the sums of every 2 entries in local buffer 2, and store them back into local buffer 2, repeat for every 4,16,32,... etc. until we have 1 sum remaining
  3. re-form the SIMD into a single float using (x.s0 + x.s1 + x.s2 + x.s3) / 4 which is then returned by the kernel
  4. on the Host (CPU) side of things we collect these floats from each core back and repeat steps 2. - 4. again (either on the CPU if the work-size is small, or re-feed them to the GPU if the work size is still large).

This should give us back the squared sum over ALL sliding windows. With O(log(N x M)) operations. Except for initial data assignment, almost all operations are performed locally and should be fast.

Note: I'm not entirely certain if buffer 1 is required, or adds any benefit, but it probably doesn't hurt.

---

As this doesn't require much adaptation from the existing summed-area table computation, the above will be used as the proposed solution.

I will attempt them in rust/opencl before numba in order to establish a baseline without any "automagic" run-time optimizations that numba perform. Ideally they give the same-ish performance, and not a huge speed-loss on the first kernel execution in the case of numba.