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 &lt; (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
) </code></pre><p style=""></p><p style=""></p><p style=""></p><p style=""></p><p style=""></p><p style=""></p><p style=""></p><p style=""></p><p style=""></p><p style=""></p><p style=""></p><p style=""></p><p style=""></p><p style=""></p><p style=""></p><p style=""></p><p style=""></p><p style=""></p>