MLSYS ENGINEERING

6.4. Op fusion

One common technique to ease the memory-bound problem is op fusion: combining multiple ops into a single pass, so data only needs to be read from and written to memory once. The following figure illustrates the idea.

Without fusion: Read Op 1 Write Read Op 2 Write With fusion: Read Op 1 Op 2 Write Time
Figure 19. Op fusion.

Here is a simple code example. We have a list of numbers stored in memory. Reading and writing each number from and to memory takes 1 second, which is simulated using time.sleep(1).

Code 17. Use sleep to simulate slow memory access.
import time

data = [float(i) for i in range(10)]

def read_from_memory(index):
    time.sleep(1)
    return data[index]

def write_to_memory(index, value):
    time.sleep(1)
    data[index] = value

You can think of this read as copying data from global memory to shared memory in a GPU, and the write as copying it back to global memory.

We want to double each number and then add 1 to each number. One way to do it is to read each number, double it, and write it back; then read each number again, add 1, and write it back. This is shown in the following code snippet.

Code 18. Inefficient memory-access example
for i in range(len(data)):
    val = read_from_memory(i)
    write_to_memory(i, val * 2)

for i in range(len(data)):
    val = read_from_memory(i)
    write_to_memory(i, val + 1)

Note that in real-world implementations, the compute should be overlapped with memory access using the pipelining we introduced in Pipelining. We did not implement that here for simplicity. We only want to show how we can reduce the number of memory accesses for the same compute. In this example, we can assume the compute is completely hidden by the memory access, so it is memory-bound.

Each iteration of the for loop takes 2 seconds (1 second for reading and 1 second for writing). We have 2 for loops with 10 iterations each. So, the total time is 40 seconds = 2 seconds * 10 iterations * 2 loops.

How do we make it faster? If we only read each number once, double it, add 1, and write it back, we can save a lot of time. This is shown in the following code snippet.

Code 19. Efficient memory-access example
for i in range(len(data)):
    val = read_from_memory(i)
    val = val * 2
    val = val + 1
    write_to_memory(i, val)

Now we only have one for loop instead of two, with 10 iterations. Each iteration still takes 2 seconds. So, the total time is 20 seconds = 2 seconds * 10 iterations. What we did was merge multiple ops into a single pass over the data, avoiding the round trip to global memory between the two ops. More generally, op fusion means not writing intermediate results back to global memory only to read them again for the next op. Fused ops keep the data in registers or shared memory and only touch global memory once. This is known as improving data locality or optimizing the memory access pattern, and it is widely used in ML compilers and GPU kernels.