Skip to content

Commit c72b497

Browse files
Oh-AfterglowOh-Afterglow
authored
[Simulator][Dataflow] Support simulating AIE kernels (#330)
1 parent 9d450ea commit c72b497

File tree

6 files changed

+154
-39
lines changed

6 files changed

+154
-39
lines changed

‎allo/backend/simulator.py

+10-1
Original file line numberDiff line numberDiff line change
@@ -515,7 +515,16 @@ def __init__(self, mod: Module, top_func_name: str, ext_libs=None):
515515
func.attributes["top"] = UnitAttr.get()
516516

517517
# Start lowering
518-
# Lower StructType first
518+
# Lower linalg for AIE
519+
pm = PassManager.parse(
520+
"builtin.module("
521+
"one-shot-bufferize,"
522+
"expand-strided-metadata,"
523+
"func.func(convert-linalg-to-affine-loops)"
524+
")"
525+
)
526+
pm.run(self.module.operation)
527+
# Lower StructType
519528
allo_d.lower_composite_type(self.module)
520529
# Reference: https://discourse.llvm.org/t/help-lowering-affine-loop-to-openmp/72441/9
521530
pm = PassManager.parse(

‎tests/dataflow/aie/test_gemm.py

+25-6
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
# Copyright Allo authors. All Rights Reserved.
22
# SPDX-License-Identifier: Apache-2.0
33

4+
import os
45
import allo.dataflow as df
56
from allo.ir.types import int16, int32
67
import numpy as np
@@ -22,13 +23,21 @@ def gemm(A: Ty[M, K], B: Ty[K, N], C: Ty[M, N]):
2223
A[pi * Mt : (pi + 1) * Mt, :], B
2324
)
2425

25-
mod = df.build(top, target="aie")
2626
A = np.random.randint(0, 64, (M, K)).astype(np.int32)
2727
B = np.random.randint(0, 64, (K, N)).astype(np.int32)
28+
if "MLIR_AIE_INSTALL_DIR" in os.environ:
29+
mod = df.build(top, target="aie")
30+
C = np.zeros((M, N)).astype(np.int32)
31+
mod(A, B, C)
32+
np.testing.assert_allclose(C, A @ B, atol=1e-5)
33+
print("PASSED!")
34+
else:
35+
print("MLIR_AIE_INSTALL_DIR unset. Skipping AIE backend test.")
36+
sim_mod = df.build(top, target="simulator")
2837
C = np.zeros((M, N)).astype(np.int32)
29-
mod(A, B, C)
38+
sim_mod(A, B, C)
3039
np.testing.assert_allclose(C, A @ B, atol=1e-5)
31-
print("PASSED!")
40+
print("Dataflow Simulator Passed!")
3241

3342

3443
def _test_gemm_2D():
@@ -46,13 +55,23 @@ def gemm(A: TyI[M, K], B: TyI[K, N], C: TyO[M, N]):
4655
A[p0 * Mt : (p0 + 1) * Mt, :], B[:, p1 * Nt : (p1 + 1) * Nt]
4756
)
4857

49-
mod = df.build(top, target="aie")
5058
A = np.random.randint(0, 64, (M, K)).astype(np.int32)
5159
B = np.random.randint(0, 64, (K, N)).astype(np.int32)
60+
61+
if "MLIR_AIE_INSTALL_DIR" in os.environ:
62+
mod = df.build(top, target="aie")
63+
C = np.zeros((M, N)).astype(np.int32)
64+
mod(A, B, C)
65+
np.testing.assert_allclose(C, A @ B, atol=1e-5)
66+
print("PASSED!")
67+
else:
68+
print("MLIR_AIE_INSTALL_DIR unset. Skipping AIE backend test.")
69+
70+
sim_mod = df.build(top, target="simulator")
5271
C = np.zeros((M, N)).astype(np.int32)
53-
mod(A, B, C)
72+
sim_mod(A, B, C)
5473
np.testing.assert_allclose(C, A @ B, atol=1e-5)
55-
print("PASSED!")
74+
print("Dataflow Simulator Passed!")
5675

5776

5877
if __name__ == "__main__":

‎tests/dataflow/aie/test_matrix.py

+25-6
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
# Copyright Allo authors. All Rights Reserved.
22
# SPDX-License-Identifier: Apache-2.0
33

4+
import os
45
import allo
56
from allo.ir.types import int32, float32
67
import allo.dataflow as df
@@ -20,12 +21,22 @@ def core(A: Ty[M, N], B: Ty[M, N]):
2021
pi = df.get_pid()
2122
B[pi * Mt : (pi + 1) * Mt, :] = allo.add(A[pi * Mt : (pi + 1) * Mt, :], 1)
2223

23-
mod = df.build(top, target="aie")
2424
A = np.random.randint(0, 100, (M, N)).astype(np.int32)
25+
26+
if "MLIR_AIE_INSTALL_DIR" in os.environ:
27+
mod = df.build(top, target="aie")
28+
B = np.zeros((M, N)).astype(np.int32)
29+
mod(A, B)
30+
np.testing.assert_allclose(B, A + 1)
31+
print("PASSED!")
32+
else:
33+
print("MLIR_AIE_INSTALL_DIR unset. Skipping AIE backend test.")
34+
35+
sim_mod = df.build(top, target="simulator")
2536
B = np.zeros((M, N)).astype(np.int32)
26-
mod(A, B)
37+
sim_mod(A, B)
2738
np.testing.assert_allclose(B, A + 1)
28-
print("PASSED!")
39+
print("Dataflow Simulator Passed!")
2940

3041

3142
def _test_matrix_matrix_add():
@@ -43,13 +54,21 @@ def core(A: Ty[M, N], B: Ty[M, N], C: Ty[M, N]):
4354
A[pi * Mt : (pi + 1) * Mt, :], B[pi * Mt : (pi + 1) * Mt, :]
4455
)
4556

46-
mod = df.build(top, target="aie")
4757
A = np.random.randint(0, 100, (M, N)).astype(np.int32)
4858
B = np.random.randint(0, 100, (M, N)).astype(np.int32)
59+
if "MLIR_AIE_INSTALL_DIR" in os.environ:
60+
mod = df.build(top, target="aie")
61+
C = np.zeros((M, N)).astype(np.int32)
62+
mod(A, B, C)
63+
np.testing.assert_allclose(C, A + B)
64+
print("PASSED!")
65+
else:
66+
print("MLIR_AIE_INSTALL_DIR unset. Skipping AIE backend test.")
67+
sim_mod = df.build(top, target="simulator")
4968
C = np.zeros((M, N)).astype(np.int32)
50-
mod(A, B, C)
69+
sim_mod(A, B, C)
5170
np.testing.assert_allclose(C, A + B)
52-
print("PASSED!")
71+
print("Dataflow Simulator Passed!")
5372

5473

5574
if __name__ == "__main__":

‎tests/dataflow/aie/test_multi_core.py

+25-6
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
# Copyright Allo authors. All Rights Reserved.
22
# SPDX-License-Identifier: Apache-2.0
33

4+
import os
45
import allo
56
from allo.ir.types import int32
67
import allo.dataflow as df
@@ -24,12 +25,22 @@ def core(A: Ty[M], B: Ty[M]):
2425
pi = df.get_pid()
2526
B[pi * Mt : (pi + 1) * Mt] = allo.add(A[pi * Mt : (pi + 1) * Mt], 1)
2627

27-
mod = df.build(top, target="aie")
2828
A = np.random.randint(0, 100, M).astype(np.int32)
29+
30+
if "MLIR_AIE_INSTALL_DIR" in os.environ:
31+
mod = df.build(top, target="aie")
32+
B = np.zeros(M).astype(np.int32)
33+
mod(A, B)
34+
np.testing.assert_allclose(B, A + 1)
35+
print("PASSED!")
36+
else:
37+
print("MLIR_AIE_INSTALL_DIR unset. Skipping AIE backend test.")
38+
39+
sim_mod = df.build(top, target="simulator")
2940
B = np.zeros(M).astype(np.int32)
30-
mod(A, B)
41+
sim_mod(A, B)
3142
np.testing.assert_allclose(B, A + 1)
32-
print("PASSED!")
43+
print("Dataflow Simulator Passed!")
3344

3445

3546
def _test_vector_vector_add():
@@ -51,13 +62,21 @@ def core(A: Ty[M], B: Ty[M], C: Ty[M]):
5162
A[pi * Mt : (pi + 1) * Mt], B[pi * Mt : (pi + 1) * Mt]
5263
)
5364

54-
mod = df.build(top, target="aie")
5565
A = np.random.randint(0, 100, M).astype(np.int32)
5666
B = np.random.randint(0, 100, M).astype(np.int32)
67+
if "MLIR_AIE_INSTALL_DIR" in os.environ:
68+
mod = df.build(top, target="aie")
69+
C = np.zeros(M).astype(np.int32)
70+
mod(A, B, C)
71+
np.testing.assert_allclose(C, A + B)
72+
print("PASSED!")
73+
else:
74+
print("MLIR_AIE_INSTALL_DIR unset. Skipping AIE backend test.")
75+
sim_mod = df.build(top, target="simulator")
5776
C = np.zeros(M).astype(np.int32)
58-
mod(A, B, C)
77+
sim_mod(A, B, C)
5978
np.testing.assert_allclose(C, A + B)
60-
print("PASSED!")
79+
print("Dataflow Simulator Passed!")
6180

6281

6382
if __name__ == "__main__":

‎tests/dataflow/aie/test_producer_consumer.py

+8-4
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
# Copyright Allo authors. All Rights Reserved.
22
# SPDX-License-Identifier: Apache-2.0
33

4+
import os
45
import allo
56
from allo.ir.types import int32
67
import allo.dataflow as df
@@ -40,10 +41,13 @@ def test_producer_consumer():
4041
np.testing.assert_allclose(B, A + 1)
4142
print("Dataflow Simulator Passed!")
4243

43-
mod = df.build(top, target="aie")
44-
mod(A, B)
45-
np.testing.assert_allclose(A + 1, B, atol=1e-5)
46-
print("Passed!")
44+
if "MLIR_AIE_INSTALL_DIR" in os.environ:
45+
mod = df.build(top, target="aie")
46+
mod(A, B)
47+
np.testing.assert_allclose(A + 1, B, atol=1e-5)
48+
print("Passed!")
49+
else:
50+
print("MLIR_AIE_INSTALL_DIR unset. Skipping AIE backend test.")
4751

4852

4953
if __name__ == "__main__":

‎tests/dataflow/aie/test_vector.py

+61-16
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
# Copyright Allo authors. All Rights Reserved.
22
# SPDX-License-Identifier: Apache-2.0
33

4+
import os
45
import allo
56
from allo.ir.types import int32, float32, bfloat16
67
import allo.dataflow as df
@@ -18,12 +19,20 @@ def top():
1819
def core(A: Ty[M], B: Ty[M]):
1920
B[:] = allo.add(A, 1)
2021

21-
mod = df.build(top, target="aie")
2222
A = np.random.randint(0, 100, M).astype(np.int32)
23+
if "MLIR_AIE_INSTALL_DIR" in os.environ:
24+
mod = df.build(top, target="aie")
25+
B = np.zeros(M).astype(np.int32)
26+
mod(A, B)
27+
np.testing.assert_allclose(B, A + 1)
28+
print("PASSED!")
29+
else:
30+
print("MLIR_AIE_INSTALL_DIR unset. Skipping AIE backend test.")
31+
sim_mod = df.build(top, target="simulator")
2332
B = np.zeros(M).astype(np.int32)
24-
mod(A, B)
33+
sim_mod(A, B)
2534
np.testing.assert_allclose(B, A + 1)
26-
print("PASSED!")
35+
print("Dataflow Simulator Passed!")
2736

2837

2938
def _test_vector_scalar_mul():
@@ -37,12 +46,20 @@ def top():
3746
def core(A: Ty[M], B: Ty[M]):
3847
B[:] = allo.mul(A, 2)
3948

40-
mod = df.build(top, target="aie")
4149
A = np.random.random(M).astype(np.float32)
50+
if "MLIR_AIE_INSTALL_DIR" in os.environ:
51+
mod = df.build(top, target="aie")
52+
B = np.zeros(M).astype(np.float32)
53+
mod(A, B)
54+
np.testing.assert_allclose(B, A * 2, rtol=1e-5)
55+
print("PASSED!")
56+
else:
57+
print("MLIR_AIE_INSTALL_DIR unset. Skipping AIE backend test.")
58+
sim_mod = df.build(top, target="simulator")
4259
B = np.zeros(M).astype(np.float32)
43-
mod(A, B)
60+
sim_mod(A, B)
4461
np.testing.assert_allclose(B, A * 2, rtol=1e-5)
45-
print("PASSED!")
62+
print("Dataflow Simulator Passed!")
4663

4764

4865
def _test_vector_vector_add():
@@ -56,13 +73,21 @@ def top():
5673
def core(A: Ty[M], B: Ty[M], C: Ty[M]):
5774
C[:] = allo.add(A, B)
5875

59-
mod = df.build(top, target="aie")
6076
A = np.random.randint(0, 100, M).astype(np.int32)
6177
B = np.random.randint(0, 100, M).astype(np.int32)
78+
if "MLIR_AIE_INSTALL_DIR" in os.environ:
79+
mod = df.build(top, target="aie")
80+
C = np.zeros(M).astype(np.int32)
81+
mod(A, B, C)
82+
np.testing.assert_allclose(C, A + B)
83+
print("PASSED!")
84+
else:
85+
print("MLIR_AIE_INSTALL_DIR unset. Skipping AIE backend test.")
86+
sim_mod = df.build(top, target="simulator")
6287
C = np.zeros(M).astype(np.int32)
63-
mod(A, B, C)
88+
sim_mod(A, B, C)
6489
np.testing.assert_allclose(C, A + B)
65-
print("PASSED!")
90+
print("Dataflow Simulator Passed!")
6691

6792

6893
def _test_vector_vector_bf16_add():
@@ -79,13 +104,23 @@ def core(A: Ty[M], B: Ty[M], C: Ty[M]):
79104

80105
A = np.random.random(M).astype(np_bfloat16)
81106
B = np.random.random(M).astype(np_bfloat16)
82-
mod = df.build(top, target="aie")
107+
if "MLIR_AIE_INSTALL_DIR" in os.environ:
108+
mod = df.build(top, target="aie")
109+
C = np.zeros(M).astype(np_bfloat16)
110+
mod(A, B, C)
111+
np.testing.assert_allclose(
112+
C.astype(np.float32), (A + B).astype(np.float32), rtol=1e-2
113+
)
114+
print("PASSED!")
115+
else:
116+
print("MLIR_AIE_INSTALL_DIR unset. Skipping AIE backend test.")
117+
sim_mod = df.build(top, target="simulator")
83118
C = np.zeros(M).astype(np_bfloat16)
84-
mod(A, B, C)
119+
sim_mod(A, B, C)
85120
np.testing.assert_allclose(
86121
C.astype(np.float32), (A + B).astype(np.float32), rtol=1e-2
87122
)
88-
print("PASSED!")
123+
print("Dataflow Simulator Passed!")
89124

90125

91126
def _test_vector_vector_mul():
@@ -99,13 +134,23 @@ def top():
99134
def core(A: Ty[M], B: Ty[M], C: Ty[M]):
100135
C[:] = allo.mul(A, B)
101136

102-
mod = df.build(top, target="aie")
103137
A = np.random.random(M).astype(np.float32)
104138
B = np.random.random(M).astype(np.float32)
139+
140+
if "MLIR_AIE_INSTALL_DIR" in os.environ:
141+
mod = df.build(top, target="aie")
142+
C = np.zeros(M).astype(np.float32)
143+
mod(A, B, C)
144+
np.testing.assert_allclose(C, A * B, rtol=1e-5)
145+
print("PASSED!")
146+
else:
147+
print("MLIR_AIE_INSTALL_DIR unset. Skipping AIE backend test.")
148+
149+
sim_mod = df.build(top, target="simulator")
105150
C = np.zeros(M).astype(np.float32)
106-
mod(A, B, C)
107-
np.testing.assert_allclose(C, A * B, rtol=1e-5)
108-
print("PASSED!")
151+
sim_mod(A, B, C)
152+
np.testing.assert_allclose(C, A * B, atol=1e-5)
153+
print("Dataflow Simulator Passed!")
109154

110155

111156
if __name__ == "__main__":

0 commit comments

Comments
 (0)