1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

// 1 dim
int threadsPerBlock = 256;
int blocksPerGrid = (N + threadsPerBlock - 1) / threadsPerBlock;
kernel<<<blocksPerGrid, threadsPerBlock>>>(...);

//2 dim
dim3 block(16, 16);   // 256 threads
dim3 grid(
    (width  + block.x - 1) / block.x,
    (height + block.y - 1) / block.y
);

kernel<<<grid, block>>>(...);

1
2
3
4
5
6

// 1 dim
int x = blockIdx.x * blockDim.x + threadIdx.x;

// 2 dim
int x = blockIdx.x * blockDim.x + threadIdx.x;
int y = blockIdx.y * blockDim.y + threadIdx.y;

1
2
3
4
5
6
7

2 dim tensor: [2, 3]
[[1, 2, 3], 
[4, 5, 6]]

stored as 1 dim tensor: [6]
- stride = [3, 2]: [1, 2, 3, 4, 5, 6] <- kernel 实际看到的
- stride = [2, 1]: [1, 4, 2, 5, 3, 6]

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22

__device__ void to_index(int ordinal, const int *shape, int *out_index, int num_dims)
{
  /**
   * Convert an ordinal to an index in the shape. Should ensure that enumerating position 0 ... size of
   * a tensor produces every index exactly once. It may not be the inverse of index_to_position.
   * Args:
   *    ordinal: ordinal position to convert
   *    shape: tensor shape
   *    out_index: return index corresponding to position
   *    num_dims: number of dimensions in the tensor
   *
   * Returns:
   *    None (Fills in out_index)
   */
  int cur_ord = ordinal;
  for (int i = num_dims - 1; i >= 0; --i)
  {
    int sh = shape[i];
    out_index[i] = cur_ord % sh;
    cur_ord /= sh;
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

__device__ int index_to_position(const int *index, const int *strides, int num_dims)
{
  /**
   * Converts a multidimensional tensor index into a single-dimensional position in storage
   * based on strides.
   * Args:
   *    index: index tuple of ints
   *    strides: tensor strides
   *    num_dims: number of dimensions in the tensor, e.g. shape/strides of [2, 3, 4] has 3 dimensions
   *
   * Returns:
   *    int - position in storage
   */
  int position = 0;
  for (int i = 0; i < num_dims; ++i)
  {
    position += index[i] * strides[i];
  }
  return position;
}

1
2
3
4
5
6
7

A = [[1, 2, 3],        # shape: (2, 3)
     [4, 5, 6]]

B = [10, 20, 30]       # shape: (3,)  → 广播成 (2, 3)

A + B = [[11, 22, 33],
         [14, 25, 36]]

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31

__device__ void broadcast_index(const int *big_index, const int *big_shape, const int *shape, int *out_index, int num_dims_big, int num_dims)
{
  /**
   * Convert a big_index into big_shape to a smaller out_index into shape following broadcasting rules.
   * In this case it may be larger or with more dimensions than the shape given.
   * Additional dimensions may need to be mapped to 0 or removed.
   *
   * Args:
   *    big_index: multidimensional index of bigger tensor
   *    big_shape: tensor shape of bigger tensor
   *    shape: tensor shape of smaller tensor
   *    nums_big_dims: number of dimensions in bigger tensor
   *    out_index: multidimensional index of smaller tensor
   *    nums_big_dims: number of dimensions in bigger tensor
   *    num_dims: number of dimensions in smaller tensor
   *
   * Returns:
   *    None (Fills in out_index)
   */
  for (int i = 0; i < num_dims; ++i)
  {
    if (shape[i] > 1)
    {
      out_index[i] = big_index[i + (num_dims_big - num_dims)];
    }
    else
    {
      out_index[i] = 0;
    }
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42

__global__ void mapKernel( float *out, int *out_shape, int *out_strides, int out_size, float *in_storage, int *in_shape, int *in_strides, int shape_size, int fn_id)
{
  /**
   * Map function. Apply a unary function to each element of the input array and store the result in the output array.
   * Optimization: Parallelize over the elements of the output array.
   *
   * You may find the following functions useful:
   * - index_to_position: converts an index to a position in a compact array
   * - to_index: converts a position to an index in a multidimensional array
   * - broadcast_index: converts an index in a smaller array to an index in a larger array
   *
   * Args:
   *  out: compact 1D array of size out_size to write the output to
   *  out_shape: shape of the output array
   *  out_strides: strides of the output array
   *  out_size: size of the output array
   *  in_storage: compact 1D array of size in_size
   *  in_shape: shape of the input array
   *  in_strides: strides of the input array
   *  shape_size: number of dimensions in the input and output arrays, assume dimensions are the same
   *  fn_id: id of the function to apply to each element of the input array
   *
   * Returns:
   *  None (Fills in out array)
   */

  int out_index[MAX_DIMS];
  int in_index[MAX_DIMS];
  
  int global_id = blockIdx.x * blockDim.x + threadIdx.x; //  计算这个 thread 的 global id，对应的就是 out 中的元素位置。注意 out 是实际存储的数组，是一维的

  // 排除多余的 threads
  if (global_id >= out_size) return;

  to_index(global_id, out_shape, out_index, shape_size);  // 将 1 dim 映射到 2dim
  broadcast_index(out_index, out_shape, in_shape, in_index, shape_size, shape_size); // 广播一下， in_index 就是我们要从输入数组中找的位置，是一个 二维的位置

  int in_pos = index_to_position(in_index, in_strides, shape_size); // 实际输入的 1 dim 中的位置
  int out_pos = index_to_position(out_index, out_strides, shape_size); // 实际输出的 1 dim 中的位置

  out[out_pos] = fn(fn_id, in_storage[in_pos]);
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50

__global__ void zipKernel( float *out, int *out_shape, int *out_strides, int out_size, int out_shape_size, float *a_storage, int *a_shape, int *a_strides, int a_shape_size, float *b_storage, int *b_shape, int *b_strides, int b_shape_size, int fn_id)
{
  /**
   * Zip function. Apply a binary function to elements of the input array a & b and store the result in the output array.
   * Optimization: Parallelize over the elements of the output array.
   *
   * You may find the following functions useful:
   * - index_to_position: converts an index to a position in a compact array
   * - to_index: converts a position to an index in a multidimensional array
   * - broadcast_index: converts an index in a smaller array to an index in a larger array
   *
   * Args:
   *  out: compact 1D array of size out_size to write the output to
   *  out_shape: shape of the output array
   *  out_strides: strides of the output array
   *  out_size: size of the output array
   *  out_shape_size: number of dimensions in the output array
   *  a_storage: compact 1D array of size in_size
   *  a_shape: shape of the input array
   *  a_strides: strides of the input array
   *  a_shape_size: number of dimensions in the input array
   *  b_storage: compact 1D array of size in_size
   *  b_shape: shape of the input array
   *  b_strides: strides of the input array
   *  b_shape_size: number of dimensions in the input array
   *  fn_id: id of the function to apply to each element of the a & b array
   *
   *
   * Returns:
   *  None (Fills in out array)
   */

  int out_index[MAX_DIMS];
  int a_index[MAX_DIMS];
  int b_index[MAX_DIMS];

  int global_id = blockIdx.x * blockDim.x + threadIdx.x;

  if (global_id >= out_size) return;

  to_index(global_id, out_shape, out_index, out_shape_size);
  broadcast_index(out_index, out_shape, a_shape, a_index, out_shape_size, a_shape_size); // 计算输入位置
  broadcast_index(out_index, out_shape, b_shape, b_index, out_shape_size, b_shape_size); // 计算输入位置

  int a_pos = index_to_position(a_index, a_strides, a_shape_size);  // 实际输入的 1 dim 中的位置
  int b_pos = index_to_position(b_index, b_strides, b_shape_size);  // 实际输入的 1 dim 中的位置
  int out_pos = index_to_position(out_index, out_strides, out_shape_size);  // 实际输出的 1 dim 中的位置

  out[out_pos] = fn(fn_id, a_storage[a_pos], b_storage[b_pos]);
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50

__global__ void reduceKernel( float *out, int *out_shape, int *out_strides, int out_size, float *a_storage, int *a_shape, int *a_strides, int reduce_dim, float reduce_value, int shape_size, int fn_id)
{
  /**
   * Reduce function. Apply a reduce function to elements of the input array a and store the result in the output array.
   * Optimization:
   * Parallelize over the reduction operation. Each kernel performs one reduction.
   * e.g. a = [[1, 2, 3], [4, 5, 6]], kernel0 computes reduce([1, 2, 3]), kernel1 computes reduce([4, 5, 6]).
   *
   * You may find the following functions useful:
   * - index_to_position: converts an index to a position in a compact array
   * - to_index: converts a position to an index in a multidimensional array
   *
   * Args:
   *  out: compact 1D array of size out_size to write the output to
   *  out_shape: shape of the output array
   *  out_strides: strides of the output array
   *  out_size: size of the output array
   *  a_storage: compact 1D array of size in_size
   *  a_shape: shape of the input array
   *  a_strides: strides of the input array
   *  reduce_dim: dimension to reduce on
   *  reduce_value: initial value for the reduction
   *  shape_size: number of dimensions in the input & output array, assert dimensions are the same
   *  fn_id: id of the reduce function, currently only support add, multiply, and max
   *
   *
   * Returns:
   *  None (Fills in out array)
   */

  __shared__ double cache[BLOCK_DIM]; // Uncomment this line if you want to use shared memory to store partial results
  int out_index[MAX_DIMS];

  int global_id = blockDim.x * blockIdx.x + threadIdx.x;

  if (global_id >= out_size) {
    return;
  }

  to_index(global_id, out_shape, out_index, shape_size);
  int out_pos = index_to_position(out_index, out_strides, shape_size);

  int reduce_size = a_shape[reduce_dim];
  for (int i = 0; i < reduce_size; ++i) {
    out_index[reduce_dim] = i;
    int a_pos = index_to_position(out_index, a_strides, shape_size);
    reduce_value = fn(fn_id, reduce_value, a_storage[a_pos]);
  }
  out[out_pos] = reduce_value;
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67

__global__ void reduceKernel( float *out, int *out_shape, int *out_strides, int out_size, float *a_storage, int *a_shape, int *a_strides, int reduce_dim, float reduce_value, int shape_size, int fn_id)
{
  /**
   * Reduce function. Apply a reduce function to elements of the input array a and store the result in the output array.
   * Optimization:
   * Parallelize over the reduction operation. Each kernel performs one reduction.
   * e.g. a = [[1, 2, 3], [4, 5, 6]], kernel0 computes reduce([1, 2, 3]), kernel1 computes reduce([4, 5, 6]).
   *
   * You may find the following functions useful:
   * - index_to_position: converts an index to a position in a compact array
   * - to_index: converts a position to an index in a multidimensional array
   *
   * Args:
   *  out: compact 1D array of size out_size to write the output to
   *  out_shape: shape of the output array
   *  out_strides: strides of the output array
   *  out_size: size of the output array
   *  a_storage: compact 1D array of size in_size
   *  a_shape: shape of the input array
   *  a_strides: strides of the input array
   *  reduce_dim: dimension to reduce on
   *  reduce_value: initial value for the reduction
   *  shape_size: number of dimensions in the input & output array, assert dimensions are the same
   *  fn_id: id of the reduce function, currently only support add, multiply, and max
   *
   *
   * Returns:
   *  None (Fills in out array)
   */

  __shared__ double cache[BLOCK_DIM]; 
  int out_index[MAX_DIMS];

  int tid = threadIdx.x;
  int out_id = blockIdx.x;
  if (out_id >= out_size) return;

  //out pos
  to_index(out_id, out_shape, out_index, shape_size);
  int out_pos = index_to_position(out_index, out_strides, shape_size);
  int reduce_size = a_shape[reduce_dim];

  // the number of threads in a block is restricted to 32. Therefore,
  // we have to add some elements before thread-level parallelization.
  // We are doing this because the tid will be exploited if a_shape[reduce_dim]
  // is greater than blockDim.
  float local_acc = reduce_value;
  for (int i = tid; i < reduce_size; i += blockDim.x) {
    out_index[reduce_dim] = i;
    int a_pos = index_to_position(out_index, a_strides, shape_size);
    local_acc = fn(fn_id, local_acc, a_storage[a_pos]);
  }

  cache[tid] = local_acc;
  __syncthreads();  // 等待所有元素在 shared memory 中加载好

  for (int s = 1; s < blockDim.x; s *= 2) {
    if (tid % (2 * s) == 0) {
      cache[tid] = fn(fn_id, cache[tid], cache[tid + s]);
    }
    __syncthreads();  // 等待这一层 (step) 的所有 thread 操作完毕
  }
  
  if (tid == 0) {
    out[out_pos] = cache[tid];  // 第一个元素就是最后的答案
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87

__global__ void MatrixMultiplyKernel( float *out, const int *out_shape, const int *out_strides, float *a_storage, const int *a_shape, const int *a_strides, float *b_storage, const int *b_shape, const int *b_strides)
{
  /**
   * Multiply two (compact) matrices into an output (also comapct) matrix. Matrix a and b are both in a batch
   * format, with shape [batch_size, m, n], [batch_size, n, p].
   * Requirements:
   * - All data must be first moved to shared memory.
   * - Only read each cell in a and b once.
   * - Only write to global memory once per kernel.
   * There is guarantee that a_shape[0] == b_shape[0], a_shape[2] == b_shape[1],
   * and out_shape[0] == a_shape[0], out_shape[1] == b_shape[1]
   *
   * Args:
   *   out: compact 1D array of size batch_size x m x p to write the output to
   *   out_shape: shape of the output array
   *   out_strides: strides of the output array
   *   a_storage: compact 1D array of size batch_size x m x n
   *   a_shape: shape of the a array
   *   a_strides: strides of the a array
   *   b_storage: compact 1D array of size batch_size x n x p
   *   b_shape: shape of the b array
   *   b_strides: strides of the b array
   *
   * Returns:
   *   None (Fills in out array)
   */

  __shared__ float a_shared[TILE][TILE];  // TILE == 32
  __shared__ float b_shared[TILE][TILE];

  int batch = blockIdx.z;
  int a_batch_stride = a_shape[0] > 1 ? a_strides[0] : 0;
  int b_batch_stride = b_shape[0] > 1 ? b_strides[0] : 0;

  int m = a_shape[1];
  int n = a_shape[2];
  int p = b_shape[2];

  int row = blockIdx.x * TILE + threadIdx.x;
  int col = blockIdx.y * TILE + threadIdx.y;

  int a_index[3] = {batch, row, 0};
  int b_index[3] = {batch, 0, col};
  int a_pos;
  int b_pos;
  float acc = 0.0f;

  int a_strides_local[3] = {a_batch_stride, a_strides[1], a_strides[2]};
  int b_strides_local[3] = {b_batch_stride, b_strides[1], b_strides[2]};

  int num_tiles = (n + TILE - 1) / TILE;

  for (int i = 0; i < num_tiles; ++i) {
    // load elements for A and B
    int a_col = i * TILE + threadIdx.y;
    a_index[2] = a_col;
    if (row < m && a_col < n) {
      a_pos = index_to_position(a_index, a_strides_local, 3);
      a_shared[threadIdx.x][threadIdx.y] = a_storage[a_pos];
    } else {
      a_shared[threadIdx.x][threadIdx.y] = 0.0f;
    }

    int b_row = i * TILE + threadIdx.x;
    b_index[1] = b_row;
    if (b_row < n && col < p) {
      b_pos = index_to_position(b_index, b_strides_local, 3);
      b_shared[threadIdx.x][threadIdx.y] = b_storage[b_pos];
    } else {
      b_shared[threadIdx.x][threadIdx.y] = 0.0f;
    }

    __syncthreads();

    // Calculation for C
    for (int k = 0; k < TILE; ++k) {
      acc += a_shared[threadIdx.x][k] * b_shared[k][threadIdx.y];
    }
    __syncthreads();
  }

  if (row < m && col < p) {
    int out_index[3] = {batch, row, col};
    int out_pos = index_to_position(out_index, out_strides, 3);
    out[out_pos] = acc;
  }
}

1

nvcc -o minitorch/cuda_kernels/combine.so --shared src/combine.cu -Xcompiler -fPIC

1

.cu 源码 → nvcc 编译 → .so 共享库 → Python ctypes 加载 → 调用 CUDA 函数

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32

import ctypes
import numpy as np

lib = ctypes.CDLL("path/to/your.so")
datatype = np.float32

# 1) 绑定签名
lib.yourKernel.argtypes = [
    np.ctypeslib.ndpointer(dtype=datatype, ndim=1, flags="C_CONTIGUOUS"),  # out_storage
    np.ctypeslib.ndpointer(dtype=np.int32, ndim=1, flags="C_CONTIGUOUS"),  # out_shape
    np.ctypeslib.ndpointer(dtype=np.int32, ndim=1, flags="C_CONTIGUOUS"),  # out_strides
    ctypes.c_int,                                                         # out_size
    np.ctypeslib.ndpointer(dtype=datatype, ndim=1, flags="C_CONTIGUOUS"),  # in_storage
    np.ctypeslib.ndpointer(dtype=np.int32, ndim=1, flags="C_CONTIGUOUS"),  # in_shape
    np.ctypeslib.ndpointer(dtype=np.int32, ndim=1, flags="C_CONTIGUOUS"),  # in_strides
    ctypes.c_int,                                                         # in_size
    # ... other args
]
lib.yourKernel.restype = None

# 2) 调用（把底层 buffer/shape/strides 塞进去）
lib.yourKernel(
    out_storage,
    out_shape.astype(np.int32),
    out_strides.astype(np.int32),
    out_size,
    in_storage,
    in_shape.astype(np.int32),
    in_strides.astype(np.int32),
    in_size,
    # ... other args
)