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: testIntegerSetCapsule
14def testIntegerSetCapsule():
15  with Context() as ctx:
16    is1 = IntegerSet.get_empty(1, 1, ctx)
17  capsule = is1._CAPIPtr
18  # CHECK: mlir.ir.IntegerSet._CAPIPtr
19  print(capsule)
20  is2 = IntegerSet._CAPICreate(capsule)
21  assert is1 == is2
22  assert is2.context is ctx
23
24run(testIntegerSetCapsule)
25
26
27# CHECK-LABEL: TEST: testIntegerSetGet
28def testIntegerSetGet():
29  with Context():
30    d0 = AffineDimExpr.get(0)
31    d1 = AffineDimExpr.get(1)
32    s0 = AffineSymbolExpr.get(0)
33    c42 = AffineConstantExpr.get(42)
34
35    # CHECK: (d0, d1)[s0] : (d0 - d1 == 0, s0 - 42 >= 0)
36    set0 = IntegerSet.get(2, 1, [d0 - d1, s0 - c42], [True, False])
37    print(set0)
38
39    # CHECK: (d0)[s0] : (1 == 0)
40    set1 = IntegerSet.get_empty(1, 1)
41    print(set1)
42
43    # CHECK: (d0)[s0, s1] : (d0 - s1 == 0, s0 - 42 >= 0)
44    set2 = set0.get_replaced([d0, AffineSymbolExpr.get(1)], [s0], 1, 2)
45    print(set2)
46
47    try:
48      IntegerSet.get(2, 1, [], [])
49    except ValueError as e:
50      # CHECK: Expected non-empty list of constraints
51      print(e)
52
53    try:
54      IntegerSet.get(2, 1, [d0 - d1], [True, False])
55    except ValueError as e:
56      # CHECK: Expected the number of constraints to match that of equality flags
57      print(e)
58
59    try:
60      IntegerSet.get(2, 1, [0], [True])
61    except RuntimeError as e:
62      # CHECK: Invalid expression when attempting to create an IntegerSet
63      print(e)
64
65    try:
66      IntegerSet.get(2, 1, [None], [True])
67    except RuntimeError as e:
68      # CHECK: Invalid expression (None?) when attempting to create an IntegerSet
69      print(e)
70
71    try:
72      set0.get_replaced([d0], [s0], 1, 1)
73    except ValueError as e:
74      # CHECK: Expected the number of dimension replacement expressions to match that of dimensions
75      print(e)
76
77    try:
78      set0.get_replaced([d0, d1], [s0, s0], 1, 1)
79    except ValueError as e:
80      # CHECK: Expected the number of symbol replacement expressions to match that of symbols
81      print(e)
82
83    try:
84      set0.get_replaced([d0, 1], [s0], 1, 1)
85    except RuntimeError as e:
86      # CHECK: Invalid expression when attempting to create an IntegerSet by replacing dimensions
87      print(e)
88
89    try:
90      set0.get_replaced([d0, d1], [None], 1, 1)
91    except RuntimeError as e:
92      # CHECK: Invalid expression (None?) when attempting to create an IntegerSet by replacing symbols
93      print(e)
94
95run(testIntegerSetGet)
96
97
98# CHECK-LABEL: TEST: testIntegerSetProperties
99def testIntegerSetProperties():
100  with Context():
101    d0 = AffineDimExpr.get(0)
102    d1 = AffineDimExpr.get(1)
103    s0 = AffineSymbolExpr.get(0)
104    c42 = AffineConstantExpr.get(42)
105
106    set0 = IntegerSet.get(2, 1, [d0 - d1, s0 - c42, s0 - d0], [True, False, False])
107    # CHECK: 2
108    print(set0.n_dims)
109    # CHECK: 1
110    print(set0.n_symbols)
111    # CHECK: 3
112    print(set0.n_inputs)
113    # CHECK: 1
114    print(set0.n_equalities)
115    # CHECK: 2
116    print(set0.n_inequalities)
117
118    # CHECK: 3
119    print(len(set0.constraints))
120
121    # CHECK-DAG: d0 - d1 == 0
122    # CHECK-DAG: s0 - 42 >= 0
123    # CHECK-DAG: -d0 + s0 >= 0
124    for cstr in set0.constraints:
125      print(cstr.expr, end='')
126      print(" == 0" if cstr.is_eq else " >= 0")
127
128run(testIntegerSetProperties)
129