「BroWithAIInfra」纯种网安老头简单学学 CMU11868

Grid                          ← 一次 kernel 启动的全部
├── Block 0                   ← 你指定有几个、每个多大
│   ├── Warp 0  (thread 0-31)     ← ⚠️ 你控制不了，硬件自动切
│   ├── Warp 1  (thread 32-63)
│   └── ...
├── Block 1
└── ...

GPU (A6000)
└── 84 × SM  (Streaming Multiprocessor)
    └── 4 × Partition          ← 每个有 1 个 warp scheduler
        └── 32 × Lane          ← 就是所谓的 "CUDA core"

__global__ void VecAddKernel(int* A, int* B, int* C, int n) {
    int i = blockDim.x * blockIdx.x + threadIdx.x;
    if (i < n) C[i] = A[i] + B[i];
}
// n = 1024
VecAddKernel<<<4, 256>>>(dA, dB, dC, 1024);

if (cond) { A(); }   // 16 lane 活跃，16 lane 空转
else      { B(); }   // 16 lane 活跃，16 lane 空转

if (threadIdx.x % 2 == 0)      { ... }   // ❌ 每个 warp 内一半一半 → 2×
if (threadIdx.x / 32 % 2 == 0) { ... }   // ✅ 整个 warp 同进同出 → 零代价

A = [[10]
     [20]]

B = [[1, 2, 3]
     [4, 5, 6]]


A + B = [[1 + 10, 2 + 10, 3 + 10],
         [4 + 20, 5 + 20, 6 + 20]]

A = [10, 20, 30]
B = [[1, 2, 3]
     [4, 5, 6]]

shape[2, 3]
shape   [3] <-- 右对齐

__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];

  /// BEGIN HW1_1
  int pos = blockIdx.x * blockDim.x + threadIdx.x;
  if (pos >= out_size)
  {
    return;
  }


  to_index(pos, out_shape, out_index, shape_size);

  // optional, as we have the same number of dimensions for in and out
  broadcast_index(out_index, out_shape, in_shape, in_index, shape_size, shape_size);

  int in_pos = index_to_position(in_index, in_strides, shape_size);
  int out_pos = index_to_position(out_index, out_strides, shape_size);

  out[out_pos] = fn(fn_id, in_storage[in_pos]);
  /// TODO
  // Hints:
  // 1. Compute the position in the output array that this thread will write to
  // 2. Convert the position to the out_index according to out_shape
  // 3. Broadcast the out_index to the in_index according to in_shape (optional in some cases)
  // 4. Calculate the position of element in in_array according to in_index and in_strides
  // 5. Calculate the position of element in out_array according to out_index and out_strides
  // 6. Apply the unary function to the input element and write the output to the out memory

  // assert(false && "Not Implemented");
  /// END HW1_1
}

__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];

  /// BEGIN HW1_2

  int pos = blockIdx.x * blockDim.x + threadIdx.x;
  if (pos >= out_size) { return; }

  to_index(pos, 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);
  int b_pos = index_to_position(b_index, b_strides, b_shape_size);
  int out_pos = index_to_position(out_index, out_strides, out_shape_size);

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

  /// TODO
  // Hints:
  // 1. Compute the position in the output array that this thread will write to
  // 2. Convert the position to the out_index according to out_shape
  // 3. Calculate the position of element in out_array according to out_index and out_strides
  // 4. Broadcast the out_index to the a_index according to a_shape
  // 5. Calculate the position of element in a_array according to a_index and a_strides
  // 6. Broadcast the out_index to the b_index according to b_shape
  // 7.Calculate the position of element in b_array according to b_index and b_strides
  // 8. Apply the binary function to the input elements in a_array & b_array and write the output to the out memory

  // assert(false && "Not Implemented");
  /// END HW1_2
}

lib.tensorZip(
                out._tensor._storage,
                out._tensor._shape.astype(np.int32),
                out._tensor._strides.astype(np.int32),
                out.size,
                len(out.shape),
                a._tensor._storage,
                a._tensor._shape.astype(np.int32),
                a._tensor._strides.astype(np.int32),
                a.size,
                len(a.shape),
                b._tensor._storage,
                b._tensor._shape.astype(np.int32),
                b._tensor._strides.astype(np.int32),
                b.size,
                len(b.shape),
                fn_id,
            )

__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];

  /// BEGIN HW1_3
  /// TODO
  int pos = blockIdx.x * blockDim.x + threadIdx.x;
  if (pos >= out_size) { return; }
  /*
  [[1, 2, 3],
   [4, 5, 6]]

  shape[2, 3] <-> shape[2, 1]
  */

  to_index(pos, out_shape, out_index, shape_size);

  int out_pos = index_to_position(out_index, out_strides, shape_size);

  float reduced_value = reduce_value;
  for (int i = 0; i < a_shape[reduce_dim]; ++i) {
    out_index[reduce_dim] = i;  // actually `a_index[reduce_dim] = i`, but we can reuse out_index since they have the same number of dimensions
    int a_pos = index_to_position(out_index, a_strides, shape_size);
    float a = a_storage[a_pos];
    reduced_value = fn(fn_id, reduced_value, a);
  }


  out[out_pos] = reduced_value;

  // 1. Define the position of the output element that this thread or this block will write to
  // 2. Convert the out_pos to the out_index according to out_shape
  // 3. Initialize the reduce_value to the output element
  // 4. Iterate over the reduce_dim dimension of the input array to compute the reduced value
  // 5. Write the reduced value to out memory

  // assert(false && "Not Implemented");
  /// END HW1_3
}

第 0 轮:  a0  a1  a2  a3  a4  a5  a6  a7
第 1 轮:  └┬┘  └┬┘  └┬┘  └┬┘      4 个线程各加一对 → 4 个部分和
          s1  s2  s3  s4
第 2 轮:  └──┬──┘  └──┬──┘          2 个线程       → 2 个部分和
             t1        t2
第 3 轮:  └─────┬─────┘             1 个线程       → 总和
                total

__shared__ double cache[BLOCK_DIM];

cache[threadIdx.x] = a_storage[pos_a];
__syncthreads();

for (int i = 1; i < a_shape[reduce_dim]; i<<=1) {
  if (threadIdx.x % (2*i) == 0 && threadIdx.x + i < a_shape[reduce_dim])
    cache[threadIdx.x] += cache[threadIdx.x + i];
  __syncthreads();
}

if (threadIdx.x == 0) out[out_pos] = cache[threadIdx.x];

[2048x262144]    naive 16.01 ms (134.2 GB/s)   tree<256> 8.69 ms (247.1 GB/s)  1.84x

__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] == a_shape[1], out_shape[2] == b_shape[2].
   *
   * 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];
  __shared__ float b_shared[TILE][TILE];

  // In each block, we will compute a batch of the output matrix
  // All the threads in the block will work together to compute this batch
  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;

  /// BEGIN HW1_4

  // for a [m, n] and b [n, p], the output is [m, p], and for each element, we'll do n multiplications
  // which needs n/TILE iterations of loading a tile of a and b into shared memory, and then computing the output tile.

  int row = blockIdx.y * TILE + threadIdx.y;
  int col = blockIdx.x * TILE + threadIdx.x;
  int out_pos = batch * out_strides[0] + row * out_strides[1] + col * out_strides[2];

  bool row_in_bounds = row < out_shape[1];
  bool col_in_bounds = col < out_shape[2];

  float sum = 0.0f;
  for (int t = 0; t < (a_shape[2] + TILE - 1) / TILE; t++) {
    int a_pos = batch * a_batch_stride + row * a_strides[1] + t * TILE * a_strides[2] + threadIdx.x;
    int b_pos = batch * b_batch_stride + t * TILE * b_strides[1] + threadIdx.y + col * b_strides[2];
    int N = min(TILE, a_shape[2] - t * TILE);

    a_shared[threadIdx.y][threadIdx.x] = (row_in_bounds && threadIdx.x < N) ? a_storage[a_pos] : 0.0f;
    b_shared[threadIdx.y][threadIdx.x] = (col_in_bounds && threadIdx.y < N) ? b_storage[b_pos] : 0.0f;
    __syncthreads();

    for (int i = 0; i < N; i++) {
      sum += a_shared[threadIdx.y][i] * b_shared[i][threadIdx.x];
    }
    __syncthreads();
  }

  if (row_in_bounds && col_in_bounds) {
    out[out_pos] = sum;
  }
  /// TODO
  // Hints:
  // 1. Compute the row and column of the output matrix this block will compute
  // 2. Compute the position in the output array that this thread will write to
  // 3. Iterate over tiles of the two input matrices, read the data into shared memory
  // 4. Synchronize to make sure the data is available to all threads
  // 5. Compute the output tile for this thread block
  // 6. Synchronize to make sure all threads are done computing the output tile for (row, col)
  // 7. Write the output to global memory

  // assert(false && "Not Implemented");
  /// END HW1_4
}