Repository navigation
Expand file tree
/
Copy pathvec_add_1d.py
More file actions
61 lines (45 loc) · 1.61 KB
/
Copy pathvec_add_1d.py
File metadata and controls
61 lines (45 loc) · 1.61 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
# Copyright (c) Tile-AI Corporation.
# Licensed under the MIT License.
import os
import tilelang
import tilelang.language as T
import torch
import torch_npu
tilelang.cache.clear_cache()
dtype = "float32"
seq_len = 4096
def vec_add(N, block_N, dtype="float32"):
n_num = N // block_N
@T.prim_func
def main(
A: T.Tensor((N), dtype),
B: T.Tensor((N), dtype),
C: T.Tensor((N), dtype),
shape: T.int32,
):
with T.Kernel(n_num, is_npu=True) as (cid, _):
A_VEC = T.alloc_ub((block_N), dtype)
B_VEC = T.alloc_ub((block_N), dtype)
C_VEC = T.alloc_ub((block_N), dtype)
t0 = cid * block_N
t0 = shape - t0
tail_size = T.min(block_N, t0)
T.copy(A[cid * block_N : cid * block_N + tail_size], A_VEC[0:tail_size])
T.copy(B[cid * block_N : cid * block_N + tail_size], B_VEC[0:tail_size])
T.vadd(A_VEC, B_VEC, C_VEC)
T.copy(C_VEC[0:tail_size], C[cid * block_N : cid * block_N + tail_size])
return main
def test_vec_add():
func = vec_add(seq_len, seq_len)
compiled_kernel = tilelang.compile(func, target="npuir")
v1 = torch.randn(size=[seq_len], dtype=eval("torch." + dtype)).npu()
v2 = torch.randn(size=[seq_len], dtype=eval("torch." + dtype)).npu()
v3 = torch.zeros(size=[seq_len], dtype=eval("torch." + dtype)).npu()
y_ref = v1 + v2
compiled_kernel(v1, v2, v3, seq_len)
print(y_ref)
print(v3)
torch.testing.assert_close(y_ref, v3, rtol=1e-2, atol=1e-2)
print("All check passed")
if __name__ == "__main__":
test_vec_add()