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