143 lines · python
1# RUN: %PYTHON %s | FileCheck %s2 3import gc4from mlir.ir import *5 6 7def run(f):8 print("\nTEST:", f.__name__)9 f()10 gc.collect()11 assert Context._get_live_count() == 012 return f13 14 15# CHECK-LABEL: TEST: testIntegerSetCapsule16@run17def testIntegerSetCapsule():18 with Context() as ctx:19 is1 = IntegerSet.get_empty(1, 1, ctx)20 capsule = is1._CAPIPtr21 # CHECK: mlir.ir.IntegerSet._CAPIPtr22 print(capsule)23 is2 = IntegerSet._CAPICreate(capsule)24 assert is1 == is225 assert is2.context is ctx26 27 28# CHECK-LABEL: TEST: testIntegerSetGet29@run30def testIntegerSetGet():31 with Context():32 d0 = AffineDimExpr.get(0)33 d1 = AffineDimExpr.get(1)34 s0 = AffineSymbolExpr.get(0)35 c42 = AffineConstantExpr.get(42)36 37 # CHECK: (d0, d1)[s0] : (d0 - d1 == 0, s0 - 42 >= 0)38 set0 = IntegerSet.get(2, 1, [d0 - d1, s0 - c42], [True, False])39 print(set0)40 41 # CHECK: (d0)[s0] : (1 == 0)42 set1 = IntegerSet.get_empty(1, 1)43 print(set1)44 45 # CHECK: (d0)[s0, s1] : (d0 - s1 == 0, s0 - 42 >= 0)46 set2 = set0.get_replaced([d0, AffineSymbolExpr.get(1)], [s0], 1, 2)47 print(set2)48 49 try:50 IntegerSet.get(2, 1, [], [])51 except ValueError as e:52 # CHECK: Expected non-empty list of constraints53 print(e)54 55 try:56 IntegerSet.get(2, 1, [d0 - d1], [True, False])57 except ValueError as e:58 # CHECK: Expected the number of constraints to match that of equality flags59 print(e)60 61 try:62 IntegerSet.get(2, 1, [0], [True])63 except RuntimeError as e:64 # CHECK: Invalid expression when attempting to create an IntegerSet65 print(e)66 67 try:68 IntegerSet.get(2, 1, [None], [True])69 except RuntimeError as e:70 # CHECK: Invalid expression (None?) when attempting to create an IntegerSet71 print(e)72 73 try:74 set0.get_replaced([d0], [s0], 1, 1)75 except ValueError as e:76 # CHECK: Expected the number of dimension replacement expressions to match that of dimensions77 print(e)78 79 try:80 set0.get_replaced([d0, d1], [s0, s0], 1, 1)81 except ValueError as e:82 # CHECK: Expected the number of symbol replacement expressions to match that of symbols83 print(e)84 85 try:86 set0.get_replaced([d0, 1], [s0], 1, 1)87 except RuntimeError as e:88 # CHECK: Invalid expression when attempting to create an IntegerSet by replacing dimensions89 print(e)90 91 try:92 set0.get_replaced([d0, d1], [None], 1, 1)93 except RuntimeError as e:94 # CHECK: Invalid expression (None?) when attempting to create an IntegerSet by replacing symbols95 print(e)96 97 98# CHECK-LABEL: TEST: testIntegerSetProperties99@run100def testIntegerSetProperties():101 with Context():102 d0 = AffineDimExpr.get(0)103 d1 = AffineDimExpr.get(1)104 s0 = AffineSymbolExpr.get(0)105 c42 = AffineConstantExpr.get(42)106 107 set0 = IntegerSet.get(2, 1, [d0 - d1, s0 - c42, s0 - d0], [True, False, False])108 # CHECK: 2109 print(set0.n_dims)110 # CHECK: 1111 print(set0.n_symbols)112 # CHECK: 3113 print(set0.n_inputs)114 # CHECK: 1115 print(set0.n_equalities)116 # CHECK: 2117 print(set0.n_inequalities)118 119 # CHECK: 3120 print(len(set0.constraints))121 122 # CHECK-DAG: d0 - d1 == 0123 # CHECK-DAG: s0 - 42 >= 0124 # CHECK-DAG: -d0 + s0 >= 0125 for cstr in set0.constraints:126 print(cstr.expr, end="")127 print(" == 0" if cstr.is_eq else " >= 0")128 129 130# TODO-LABEL: TEST: testHash131@run132def testHash():133 with Context():134 d0 = AffineDimExpr.get(0)135 d1 = AffineDimExpr.get(1)136 set = IntegerSet.get(2, 0, [d0 + d1], [True])137 138 assert hash(set) == hash(IntegerSet.get(2, 0, [d0 + d1], [True]))139 140 dictionary = dict()141 dictionary[set] = 42142 assert set in dictionary143