Matrix Add & 2-D Indexing

Two new habits: index mapping and 2-D launch geometry

Vector add launched a 1-D grid. Matrix add launches a 2-D grid because the data is 2-D. The launch shape mirrors the data shape, and the indexing inside the kernel maps thread coordinates to data coordinates.

The mapping has two parts:

Thread coords → data coords: row = blockIdx.y * blockDim.y + threadIdx.y, col = blockIdx.x * blockDim.x + threadIdx.x.
Data coords → flat memory offset: idx = row * ld + col for row-major; reverse for column-major.

Block size choice: 16×16 = 256 threads is a robust default. It's a multiple of the warp size (32), gives the scheduler ~8 warps for latency hiding, and balances occupancy vs. register pressure. 32×32 = 1024 also works but can hit register limits in heavier kernels.

Code

Matrix add — 2-D launch + 2-D indexing (CUDA)·cuda

__global__ void mat_add(const float* A, const float* B, float* C,
                        int rows, int cols, int ld) {
    int row = blockIdx.y * blockDim.y + threadIdx.y;
    int col = blockIdx.x * blockDim.x + threadIdx.x;
    if (row < rows && col < cols) {
        size_t idx = (size_t)row * ld + col;     // row-major
        C[idx] = A[idx] + B[idx];
    }
}

// Launch:
//   dim3 block(16, 16);                      // 16×16 = 256 threads
//   dim3 grid((cols + 15)/16, (rows + 15)/16);
//   mat_add<<<grid, block>>>(dA, dB, dC, rows, cols, ld);

Same kernel in Metal·metal

#include <metal_stdlib>
using namespace metal;

kernel void mat_add(const device float *A [[buffer(0)]],
                    const device float *B [[buffer(1)]],
                    device       float *C [[buffer(2)]],
                    constant     uint2 &dims [[buffer(3)]],   // (rows, cols)
                    constant     uint  &ld   [[buffer(4)]],
                    uint2 gid [[thread_position_in_grid]])
{
    uint row = gid.y, col = gid.x;
    if (row < dims.x && col < dims.y) {
        uint idx = row * ld + col;
        C[idx] = A[idx] + B[idx];
    }
}

// Host: encoder.dispatchThreads(
//   MTLSize(width: cols, height: rows, depth: 1),
//   threadsPerThreadgroup: MTLSize(width: 16, height: 16, depth: 1))

Exercise

Write the matrix add kernel above for a 4096×4096 matrix and verify the result against numpy(A + B). Then deliberately swap the indices: idx = col * ld + row (treating the buffer as column-major while reading row-major). The math gives the wrong answer at most positions — useful to feel why layout has to be respected, not assumed.

Two new habits: index mapping and 2-D launch geometry

Code

External links

Exercise

Progress

Comments 0