はじめてのにき

_ triton

https://gigazine.net/news/20210730-triton/

25行以下はすごいなと

https://triton-lang.org/getting-started/tutorials/03-matrix-multiplication.html#sphx-glr-getting-started-tutorials-03-matrix-multiplication-py

を見ると…。うーん？

   # matrix multiplication
   pid = tl.program_id(0)
   grid_m = (M + BLOCK_M - 1) // BLOCK_M
   grid_n = (N + BLOCK_N - 1) // BLOCK_N
   # re-order program ID for better L2 performance
   width = GROUP_M * grid_n
   group_id = pid // width
   group_size = min(grid_m - group_id * GROUP_M, GROUP_M)
   pid_m = group_id * GROUP_M + (pid % group_size)
   pid_n = (pid % width) // (group_size)
   # do matrix multiplication
   rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
   rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
   rk = tl.arange(0, BLOCK_K)
   A = A + (rm[:, None] * stride_am + rk[None, :] * stride_ak)
   B = B + (rk[:, None] * stride_bk + rn[None, :] * stride_bn)
   acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
   for k in range(K, 0, -BLOCK_K):
       a = tl.load(A)
       b = tl.load(B)
       acc += tl.dot(a, b)
       A += BLOCK_K * stride_ak
       B += BLOCK_K * stride_bk

をさしている感じかな？

(20:12)