Skip to content

Commit

Permalink
Fix several tests
Browse files Browse the repository at this point in the history
  • Loading branch information
yamaguchi1024 committed Jan 14, 2025
1 parent 547a973 commit 20fb4b9
Show file tree
Hide file tree
Showing 14 changed files with 575 additions and 145 deletions.
2 changes: 1 addition & 1 deletion src/exo/rewrite/dataflow_pprint.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,7 @@ def get_name(self, nm):
if resolved := self.env.get(nm):
return resolved

candidate = repr(nm)
candidate = str(nm)
num = self.names.get(candidate, 1)
while candidate in self.names:
candidate = f"{nm}_{num}"
Expand Down
40 changes: 40 additions & 0 deletions tests/golden/test_dataflow/test_abs_pprint.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1,40 @@
var: y[i,i]
term: - (i-d)
- ((i-d)+1)
- x[i,d]
- y[i,i]
- x[i,d]
- y[i,i]
- x[i,d]
src: \x. \y
- (i-d)%3
- (i-d)
- ((i-d)+1)
- x_1[i,d]
- y_1[i,i]
- x_1[i,d]
- y_1[i,i]
- x_1[i,d]
- 4.0
new tree: \x. \y
- (i-d)%3
- (i-d)
- ((i-d)+1)
- x_1[i,d]
- (i-d)
- ((i-d)+1)
- x_1[i,d]
- y_1[i,i]
- x_1[i,d]
- y_1[i,i]
- x_1[i,d]
- x_1[i,d]
- (i-d)
- ((i-d)+1)
- x_1[i,d]
- y_1[i,i]
- x_1[i,d]
- y_1[i,i]
- x_1[i,d]
- x_1[i,d]
- 4.0
10 changes: 6 additions & 4 deletions tests/golden/test_dataflow/test_print.txt
Original file line number Diff line number Diff line change
@@ -1,8 +1,10 @@
def foo(x : R[3], y : R[3], z : R):
----------------------- {}
z_1 = \phi(True ? x[0] * y[2] : z)
----------------------- {}
z_2 = \phi(True ? 4.2 : z_1)
----------------------- {}
z_3 = \phi(True ? 2.0 : z_2)
----------------------- {}
------------------------ z_1 :
- ⊤
z_2 :
- 4.2
z_3 :
- 2.0
60 changes: 60 additions & 0 deletions tests/golden/test_dataflow/test_print_0.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1,60 @@
def foo(z : R[3]):
z_1 =\d0 \phi(0 == d0 ? 1.0 : z[d0])
for i in seq(0, 3):
z_2 =\i \d0 \phi(i == 0 ? z_1[d0] : z_3[i - 1, d0]) # LoopStart
z_3 =\i \d0 \phi(i == d0 ? 3.0 : z_2[i, d0])
------------------------ z_3 : \i. \d0
- (i-d0)
- z_1[d0]
- 3.0
- ((i-1)-d0)
- z_1[d0]
- 3.0
- 3.0
z_2 : \i. \d0
- i
- ⊥
- z_1[d0]
- ((i-1)-d0)
- z_1[d0]
- 3.0
- 3.0
z_4 =\d0 \phi(3 > 0 ? z_3[3 - 1, d0] : z_1[d0]) # LoopExit
z_5 =\d0 \phi(2 == d0 ? 2.0 : z_4[d0])
------------------------ z_1 : \d0
- d0
- z[d0]
- 1.0
- z[d0]
z_3 : \i. \d0
- (i-d0)
- z_1[d0]
- 3.0
- ((i-1)-d0)
- z_1[d0]
- 3.0
- 3.0
z_2 : \i. \d0
- i
- ⊥
- z_1[d0]
- ((i-1)-d0)
- z_1[d0]
- 3.0
- 3.0
z_4 : \d0
- ((3-1)-d0)
- z[d0]
- 3.0
- (((3-1)-1)-d0)
- z[d0]
- 3.0
- 3.0
z_5 : \d0
- (2-d0)
- z[d0]
- 2.0
- (((3-1)-1)-d0)
- z[d0]
- 3.0
- 3.0
7 changes: 4 additions & 3 deletions tests/golden/test_dataflow/test_simple.txt
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
def foo(z : R, x : R[3]):
----------------------- {}
z_1 = \phi(True ? 4.2 : z)
----------------------- {}
z_2 = \phi(True ? 2.0 : z_1)
----------------------- {}
------------------------ z_1 :
- 4.2
z_2 :
- 2.0
42 changes: 32 additions & 10 deletions tests/golden/test_dataflow/test_simple2.txt
Original file line number Diff line number Diff line change
@@ -1,12 +1,34 @@
def foo(z : R, x : R[3]):
----------------------- {}
z_1 = \phi(True ? 4.2 : z)
----------------------- {}
x_1 =\d0 \phi(d0 == 0 and True ? 2.0 : x[d0])
----------------------- {}
x_2 =\d0 \phi(d0 == 1 and True ? 3.0 : x_1[d0])
----------------------- {}
x_3 =\d0 \phi(d0 == 2 and True ? 5.0 : x_2[d0])
----------------------- {}
x_4 =\d0 \phi(d0 == 0 and True ? 12.0 : x_3[d0])
----------------------- {}
x_1 =\d0 \phi(0 == d0 ? 2.0 : x[d0])
x_2 =\d0 \phi(1 == d0 ? 3.0 : x_1[d0])
x_3 =\d0 \phi(2 == d0 ? 5.0 : x_2[d0])
x_4 =\d0 \phi(0 == d0 ? 12.0 : x_3[d0])
------------------------ z_1 :
- 4.2
x_1 : \d0
- d0
- x[d0]
- 2.0
- x[d0]
x_2 : \d0
- (1-d0)
- x[d0]
- 3.0
- 2.0
x_3 : \d0
- (2-d0)
- x[d0]
- 5.0
- (1-d0)
- x[d0]
- 3.0
- 2.0
x_4 : \d0
- d0
- x[d0]
- 12.0
- (2-d0)
- x[d0]
- 5.0
- 3.0
88 changes: 73 additions & 15 deletions tests/golden/test_dataflow/test_simple3.txt
Original file line number Diff line number Diff line change
@@ -1,21 +1,79 @@
def foo(z : R, n : size, x : R[3]):
----------------------- {}
z_1 = \phi(True ? 4.2 : z)
----------------------- {}
x_1 =\d0 \phi(d0 == 0 and True ? 2.0 : x[d0])
----------------------- {}
x_1 =\d0 \phi(0 == d0 ? 2.0 : x[d0])
if n < 3:
--------------------- {}
x_2 =\d0 \phi(d0 == n and True ? 3.0 : x_1[d0])
--------------------- {}
x_2 =\d0 \phi(n == d0 ? 3.0 : x_1[d0])
------------------------ z_1 : \n
- 4.2
x_1 : \n. \d0
- d0
- x[d0]
- 2.0
- x[d0]
x_2 : \n. \d0
- (n-d0)
- x[d0]
- 3.0
- d0
- x[d0]
- 2.0
- x[d0]
else:
--------------------- {}
----------------------- {}
------------------------ z_1 : \n
- 4.2
x_1 : \n. \d0
- d0
- x[d0]
- 2.0
- x[d0]
z_2 = \phi(n < 3 ? z_1 : z_1) # IfJoin
----------------------- {}
x_3 =\d0 \phi(n < 3 ? x_2[d0] : x_1[d0]) # IfJoin
----------------------- {}
x_4 =\d0 \phi(d0 == 2 and True ? 5.0 : x_3[d0])
----------------------- {}
x_5 =\d0 \phi(d0 == 0 and True ? 12.0 : x_4[d0])
----------------------- {}
x_4 =\d0 \phi(2 == d0 ? 5.0 : x_3[d0])
x_5 =\d0 \phi(0 == d0 ? 12.0 : x_4[d0])
------------------------ z_1 : \n
- 4.2
x_1 : \n. \d0
- d0
- x[d0]
- 2.0
- x[d0]
x_2 : \n. \d0
- (n-d0)
- x[d0]
- 3.0
- d0
- x[d0]
- 2.0
- x[d0]
z_2 : \n
- 4.2
x_3 : \n. \d0
- (n-3)
- 3.0
- d0
- x[d0]
- 2.0
- x[d0]
- d0
- x[d0]
- 2.0
- x[d0]
x_4 : \n. \d0
- (2-d0)
- (d0-3)
- x[d0]
- x[d0]
- x[d0]
- 5.0
- 3.0
x_5 : \n. \d0
- d0
- x[d0]
- 12.0
- (2-d0)
- (d0-3)
- x[d0]
- x[d0]
- x[d0]
- 5.0
- 3.0
10 changes: 7 additions & 3 deletions tests/golden/test_dataflow/test_simple_stmts.txt
Original file line number Diff line number Diff line change
@@ -1,6 +1,10 @@
def foo(z : R, x : R[3]):
----------------------- {}
z_1 = \phi(True ? 4.2 : z)
----------------------- {}
z_2 = \phi(True ? 2.0 : z_1)
----------------------- {}[DataflowIR.Assign(lhs=z_26, iters=[], dims=[], cond=DataflowIR.Const(val=True, type=DataflowIR.Bool(), srcinfo=<exo.prelude.SrcInfo object at 0x7f4c89ce85e0>), body=DataflowIR.Const(val=4.2, type=DataflowIR.Num(), srcinfo=<exo.prelude.SrcInfo object at 0x7f4c874ccc10>), orelse=DataflowIR.Read(name=z_23, idx=[], type=DataflowIR.Num(), srcinfo=<exo.prelude.SrcInfo object at 0x7f4c89ce85e0>), srcinfo=<exo.prelude.SrcInfo object at 0x7f4c874cc8e0>)][DataflowIR.Assign(lhs=z_27, iters=[], dims=[], cond=DataflowIR.Const(val=True, type=DataflowIR.Bool(), srcinfo=<exo.prelude.SrcInfo object at 0x7f4c89ce85e0>), body=DataflowIR.Const(val=2.0, type=DataflowIR.Num(), srcinfo=<exo.prelude.SrcInfo object at 0x7f4c874cc850>), orelse=DataflowIR.Read(name=z_26, idx=[], type=DataflowIR.Num(), srcinfo=<exo.prelude.SrcInfo object at 0x7f4c89ce85e0>), srcinfo=<exo.prelude.SrcInfo object at 0x7f4c874cccd0>)]
------------------------ z_1 :
- 4.2
z_2 :
- 2.0

z = \phi(True ? 4.2 : z_1)
z = \phi(True ? 2.0 : z_1)
Loading

0 comments on commit 20fb4b9

Please sign in to comment.