1# RUN: %PYTHON %s | FileCheck %s
2
3import gc
4from mlir.ir import *
5
6def run(f):
7  print("\nTEST:", f.__name__)
8  f()
9  gc.collect()
10  assert Context._get_live_count() == 0
11
12
13# CHECK-LABEL: TEST: testAffineExprCapsule
14def testAffineExprCapsule():
15  with Context() as ctx:
16    affine_expr = AffineExpr.get_constant(42)
17
18  affine_expr_capsule = affine_expr._CAPIPtr
19  # CHECK: capsule object
20  # CHECK: mlir.ir.AffineExpr._CAPIPtr
21  print(affine_expr_capsule)
22
23  affine_expr_2 = AffineExpr._CAPICreate(affine_expr_capsule)
24  assert affine_expr == affine_expr_2
25  assert affine_expr_2.context == ctx
26
27run(testAffineExprCapsule)
28
29
30# CHECK-LABEL: TEST: testAffineExprEq
31def testAffineExprEq():
32  with Context():
33    a1 = AffineExpr.get_constant(42)
34    a2 = AffineExpr.get_constant(42)
35    a3 = AffineExpr.get_constant(43)
36    # CHECK: True
37    print(a1 == a1)
38    # CHECK: True
39    print(a1 == a2)
40    # CHECK: False
41    print(a1 == a3)
42    # CHECK: False
43    print(a1 == None)
44    # CHECK: False
45    print(a1 == "foo")
46
47run(testAffineExprEq)
48
49
50# CHECK-LABEL: TEST: testAffineExprContext
51def testAffineExprContext():
52  with Context():
53    a1 = AffineExpr.get_constant(42)
54  with Context():
55    a2 = AffineExpr.get_constant(42)
56
57  # CHECK: False
58  print(a1 == a2)
59
60run(testAffineExprContext)
61
62
63# CHECK-LABEL: TEST: testAffineExprConstant
64def testAffineExprConstant():
65  with Context():
66    a1 = AffineExpr.get_constant(42)
67    # CHECK: 42
68    print(a1.value)
69    # CHECK: 42
70    print(a1)
71
72    a2 = AffineConstantExpr.get(42)
73    # CHECK: 42
74    print(a2.value)
75    # CHECK: 42
76    print(a2)
77
78    assert a1 == a2
79
80run(testAffineExprConstant)
81
82
83# CHECK-LABEL: TEST: testAffineExprDim
84def testAffineExprDim():
85  with Context():
86    d1 = AffineExpr.get_dim(1)
87    d11 = AffineDimExpr.get(1)
88    d2 = AffineDimExpr.get(2)
89
90    # CHECK: 1
91    print(d1.position)
92    # CHECK: d1
93    print(d1)
94
95    # CHECK: 2
96    print(d2.position)
97    # CHECK: d2
98    print(d2)
99
100    assert d1 == d11
101    assert d1 != d2
102
103run(testAffineExprDim)
104
105
106# CHECK-LABEL: TEST: testAffineExprSymbol
107def testAffineExprSymbol():
108  with Context():
109    s1 = AffineExpr.get_symbol(1)
110    s11 = AffineSymbolExpr.get(1)
111    s2 = AffineSymbolExpr.get(2)
112
113    # CHECK: 1
114    print(s1.position)
115    # CHECK: s1
116    print(s1)
117
118    # CHECK: 2
119    print(s2.position)
120    # CHEKC: s2
121    print(s2)
122
123    assert s1 == s11
124    assert s1 != s2
125
126run(testAffineExprSymbol)
127
128
129# CHECK-LABEL: TEST: testAffineAddExpr
130def testAffineAddExpr():
131  with Context():
132    d1 = AffineDimExpr.get(1)
133    d2 = AffineDimExpr.get(2)
134    d12 = AffineExpr.get_add(d1, d2)
135    # CHECK: d1 + d2
136    print(d12)
137
138    d12op = d1 + d2
139    # CHECK: d1 + d2
140    print(d12op)
141
142    assert d12 == d12op
143    assert d12.lhs == d1
144    assert d12.rhs == d2
145
146run(testAffineAddExpr)
147
148
149# CHECK-LABEL: TEST: testAffineMulExpr
150def testAffineMulExpr():
151  with Context():
152    d1 = AffineDimExpr.get(1)
153    c2 = AffineConstantExpr.get(2)
154    expr = AffineExpr.get_mul(d1, c2)
155    # CHECK: d1 * 2
156    print(expr)
157
158    # CHECK: d1 * 2
159    op = d1 * c2
160    print(op)
161
162    assert expr == op
163    assert expr.lhs == d1
164    assert expr.rhs == c2
165
166run(testAffineMulExpr)
167
168
169# CHECK-LABEL: TEST: testAffineModExpr
170def testAffineModExpr():
171  with Context():
172    d1 = AffineDimExpr.get(1)
173    c2 = AffineConstantExpr.get(2)
174    expr = AffineExpr.get_mod(d1, c2)
175    # CHECK: d1 mod 2
176    print(expr)
177
178    # CHECK: d1 mod 2
179    op = d1 % c2
180    print(op)
181
182    assert expr == op
183    assert expr.lhs == d1
184    assert expr.rhs == c2
185
186run(testAffineModExpr)
187
188
189# CHECK-LABEL: TEST: testAffineFloorDivExpr
190def testAffineFloorDivExpr():
191  with Context():
192    d1 = AffineDimExpr.get(1)
193    c2 = AffineConstantExpr.get(2)
194    expr = AffineExpr.get_floor_div(d1, c2)
195    # CHECK: d1 floordiv 2
196    print(expr)
197
198    assert expr.lhs == d1
199    assert expr.rhs == c2
200
201run(testAffineFloorDivExpr)
202
203
204# CHECK-LABEL: TEST: testAffineCeilDivExpr
205def testAffineCeilDivExpr():
206  with Context():
207    d1 = AffineDimExpr.get(1)
208    c2 = AffineConstantExpr.get(2)
209    expr = AffineExpr.get_ceil_div(d1, c2)
210    # CHECK: d1 ceildiv 2
211    print(expr)
212
213    assert expr.lhs == d1
214    assert expr.rhs == c2
215
216run(testAffineCeilDivExpr)
217
218
219# CHECK-LABEL: TEST: testAffineExprSub
220def testAffineExprSub():
221  with Context():
222    d1 = AffineDimExpr.get(1)
223    d2 = AffineDimExpr.get(2)
224    expr = d1 - d2
225    # CHECK: d1 - d2
226    print(expr)
227
228    assert expr.lhs == d1
229    rhs = AffineMulExpr(expr.rhs)
230    # CHECK: d2
231    print(rhs.lhs)
232    # CHECK: -1
233    print(rhs.rhs)
234
235run(testAffineExprSub)
236
237
238def testClassHierarchy():
239  with Context():
240    d1 = AffineDimExpr.get(1)
241    c2 = AffineConstantExpr.get(2)
242    add = AffineAddExpr.get(d1, c2)
243    mul = AffineMulExpr.get(d1, c2)
244    mod = AffineModExpr.get(d1, c2)
245    floor_div = AffineFloorDivExpr.get(d1, c2)
246    ceil_div = AffineCeilDivExpr.get(d1, c2)
247
248    # CHECK: False
249    print(isinstance(d1, AffineBinaryExpr))
250    # CHECK: False
251    print(isinstance(c2, AffineBinaryExpr))
252    # CHECK: True
253    print(isinstance(add, AffineBinaryExpr))
254    # CHECK: True
255    print(isinstance(mul, AffineBinaryExpr))
256    # CHECK: True
257    print(isinstance(mod, AffineBinaryExpr))
258    # CHECK: True
259    print(isinstance(floor_div, AffineBinaryExpr))
260    # CHECK: True
261    print(isinstance(ceil_div, AffineBinaryExpr))
262
263    try:
264      AffineBinaryExpr(d1)
265    except ValueError as e:
266      # CHECK: Cannot cast affine expression to AffineBinaryExpr
267      print(e)
268
269    try:
270      AffineBinaryExpr(c2)
271    except ValueError as e:
272      # CHECK: Cannot cast affine expression to AffineBinaryExpr
273      print(e)
274
275run(testClassHierarchy)
276