Peeking into Broadcasting

So the other day, as I was working in Pytorch and Numpy, making shapes of objects into their necessary form, I wondered about how broadcasting worked under the hood. I knew it was an efficient implementation of matrix operations but I did not know how it worked.

So broadcasting in its simplest is transforming matrix shapes so that a mathematical operation can happen without duplicating data in memory. And you may wonder why we need to transform the shapes in the first place since we do not do that by hand in linear algebra class. That is because these operations are parallelized on GPUs and Single Instruction Multiple Data (SIMD) on CPUs. What’s that look like concretely?

Well, let’s say you have A = [[10,20,30], [40,50,60]] and B = [1,2,3]. A’s shape is (2,3) and B’s shape is (3,).

It would be a lot easier for CPUs and GPUs to process A@B.T if B = [[1,2,3],[1,2,3]] i.e. B’s shape is (2,3) as well. Because now one worker can work on the elemental multiplication of [10,20,30] and [1,2,3] to give 140. And another worker can work on [40,50,60] and [1,2,3] to get 320. Combining the result would then be [140,320].

And without broadcasting, the above example would result in duplicate data because now B takes up twice as much memory. And while this might not be much in a small vector, this quickly becomes computational in a 1MB image with 1 million pixels.

So how does broadcasting prevent duplicate data in memory?

It uses the concept of strides which tells the computer how many elements to jump forward in each direction. So B’s stride is (3,1) because it takes 3 steps to get to the next row and 1 step to get to the next column. And it then calculates a new stride by zeroing out the dimension you want duplicated so in this case it is the row hence (0,1).

So with this new stride, it calculates the index to look at for each position of B without requiring duplicated memory. For instance:

0 × 0 + 0 × 1 = 0 → [[1 _ _], [_ _ _]]

0 × 0 + 1 × 1 = 1 → [[_ 2 _], [_ _ _]]

0 × 0 + 2 × 1 = 2 → [[_ _ 3], [_ _ _]]

1 × 0 + 0 × 1 = 0 → [[_ _ _], [1 _ _]]

1 × 0 + 1 × 1 = 1 → [[_ _ _], [_ 2 _]]

1 × 0 + 2 × 1 = 2 → [[_ _ _], [_ _ 3]]