mirror of
https://github.com/ROCm/ROCm.git
synced 2026-02-21 03:00:39 -05:00
ba91f39dbf
This PR; - Fixes syntax errors like `.type values: dict[str, Callable[[list[Any]], Any]]` to `:type values: dict[str, Callable[[list[Any]], Any]]`, - Fixes typos, - Fixes formatting like `k ++` to ` k++`, - Increases consistency (e.g. by transforming the minority `cd dir/` to the majority `cd dir`).
19 lines
531 B
Python
19 lines
531 B
Python
import triton
|
|
import triton.language as tl
|
|
|
|
|
|
# triton kernel
|
|
@triton.jit
|
|
def kernel(X, stride_xm,
|
|
Z, stride_zn,
|
|
BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr):
|
|
off_m = tl.arange(0, BLOCK_M)
|
|
off_n = tl.arange(0, BLOCK_N)
|
|
Xs = X + off_m[:, None] * stride_xm + off_n[None, :] * 1
|
|
Zs = Z + off_m[:, None] * 1 + off_n[None, :] * stride_zn
|
|
tl.store(Zs, tl.load(Xs))
|
|
|
|
|
|
ret = triton.compile(kernel, signature="*fp32,i32,*fp32,i32", constants={"BLOCK_M": 64, "BLOCK_N": 64})
|
|
print(ret.asm["ttgir"])
|