Vector addition

原题：编写一个在 GPU 上执行 32 位浮点数向量逐元素相加的程序。该程序应接受两个等长输入向量，并产生一个包含它们和的输出向量。

exp1：

Input:  A = [1.0, 2.0, 3.0, 4.0]
        B = [5.0, 6.0, 7.0, 8.0]
Output: C = [6.0, 8.0, 10.0, 12.0]

思路：

计算block_size，load，python相加，store
根据n_elements的大小自定义block_size大小

import torch
import triton
import triton.language as tl

def get_optimal_block_size(n_elements: int) -> int:
    if n_elements < 1024: return 256
    elif n_elements < 10000: return 1024
    elif n_elements < 100000: return 2048
    elif n_elements < 1000000: return 4096
    else: return 8192

@triton.jit
def vector_add_kernel(a, b, c, n_elements, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(axis=0)
    block_start = pid * BLOCK_SIZE
    offset = block_start + tl.arange(0, BLOCK_SIZE)
    mask = offset < n_elements

    a_nums = tl.load(a + offset, mask = mask)
    b_nums = tl.load(b + offset, mask = mask)

    c_nums = a_nums + b_nums

    tl.store(c + offset, c_nums, mask = mask)

   
# a, b, c are tensors on the GPU
def solve(a: torch.Tensor, b: torch.Tensor, c: torch.Tensor, N: int):    
    BLOCK_SIZE = get_optimal_block_size(a.numel())
    grid = (triton.cdiv(N, BLOCK_SIZE),)
    vector_add_kernel[grid](a, b, c, N, BLOCK_SIZE)

结果为：

A100: 0.22ms，67.1th percentile
好像用了动态的block_size区别不大，这里挖个坑，回头看看怎么回事

Reverse Array

原题：编写一个程序，原地反转一个 32 位浮点数数组。程序应原地反转 input 。

exp1：

Input: [1.0, 2.0, 3.0, 4.0]
Output: [4.0, 3.0, 2.0, 1.0]

思路：

反转整个数组，假设有1024个数字
- 0就跟1023互换，1就跟1022互换，也就是i和N-i-1互换
- threads的数量也只需要一半，从0到511一共512个threads
- 故此block只需要覆盖一半的数字。对于单数个数字来说，33个数字的反转，只需要覆盖前16个数字即可，对于偶数个数字，34个数字需要覆盖前17个数字，故此这里需要用个整除。

import torch
import triton
import triton.language as tl

def get_optimal_block_size(n_elements: int) -> int:
    if n_elements < 1024: return 256
    elif n_elements < 10000: return 512
    elif n_elements < 1000000: return 1024
    elif n_elements < 10000000: return 2048
    else: return 4096


@triton.jit
def reverse_kernel(
    input,
    N,
    BLOCK_SIZE: tl.constexpr
):
    # 常规两件套
    pid = tl.program_id(axis = 0)
    block_start = pid * BLOCK_SIZE
    
    # 计算偏移量
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < (N // 2)
    reverse_offsets = N - 1 - offsets

    # 拿取左边和右边的数组
    left_nums = tl.load(input + offsets, mask = mask, other = 0.0)
    right_nums = tl.load(input + reverse_offsets, mask = mask, other = 0.0)

    # 反过来存放
    tl.store(input + reverse_offsets, left_nums, mask = mask)
    tl.store(input + offsets, right_nums, mask = mask)


# input is a tensor on the GPU
def solve(input: torch.Tensor, N: int):
    BLOCK_SIZE = get_optimal_block_size(N)
    n_blocks = triton.cdiv(N // 2, BLOCK_SIZE)
    grid = (n_blocks,)
    
    reverse_kernel[grid](
        input,
        N,
        BLOCK_SIZE
    )

Vector addition#

Reverse Array#

Vector addition

Reverse Array