427 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: testAffineExprCapsule16@run17def testAffineExprCapsule():18 with Context() as ctx:19 affine_expr = AffineExpr.get_constant(42)20 21 affine_expr_capsule = affine_expr._CAPIPtr22 # CHECK: capsule object23 # CHECK: mlir.ir.AffineExpr._CAPIPtr24 print(affine_expr_capsule)25 26 affine_expr_2 = AffineExpr._CAPICreate(affine_expr_capsule)27 assert affine_expr == affine_expr_228 assert affine_expr_2.context == ctx29 30 31# CHECK-LABEL: TEST: testAffineExprEq32@run33def testAffineExprEq():34 with Context():35 a1 = AffineExpr.get_constant(42)36 a2 = AffineExpr.get_constant(42)37 a3 = AffineExpr.get_constant(43)38 # CHECK: True39 print(a1 == a1)40 # CHECK: True41 print(a1 == a2)42 # CHECK: False43 print(a1 == a3)44 # CHECK: False45 print(a1 is None)46 # CHECK: False47 print(a1 == "foo")48 49 50# CHECK-LABEL: TEST: testAffineExprContext51@run52def testAffineExprContext():53 with Context():54 a1 = AffineExpr.get_constant(42)55 with Context():56 a2 = AffineExpr.get_constant(42)57 58 # CHECK: False59 print(a1 == a2)60 61 62run(testAffineExprContext)63 64 65# CHECK-LABEL: TEST: testAffineExprConstant66@run67def testAffineExprConstant():68 with Context():69 a1 = AffineExpr.get_constant(42)70 # CHECK: 4271 print(a1.value)72 # CHECK: 4273 print(a1)74 75 a2 = AffineConstantExpr.get(42)76 # CHECK: 4277 print(a2.value)78 # CHECK: 4279 print(a2)80 81 assert a1 == a282 83 84# CHECK-LABEL: TEST: testAffineExprDim85@run86def testAffineExprDim():87 with Context():88 d1 = AffineExpr.get_dim(1)89 d11 = AffineDimExpr.get(1)90 d2 = AffineDimExpr.get(2)91 92 # CHECK: 193 print(d1.position)94 # CHECK: d195 print(d1)96 97 # CHECK: 298 print(d2.position)99 # CHECK: d2100 print(d2)101 102 assert d1 == d11103 assert d1 != d2104 105 106# CHECK-LABEL: TEST: testAffineExprSymbol107@run108def testAffineExprSymbol():109 with Context():110 s1 = AffineExpr.get_symbol(1)111 s11 = AffineSymbolExpr.get(1)112 s2 = AffineSymbolExpr.get(2)113 114 # CHECK: 1115 print(s1.position)116 # CHECK: s1117 print(s1)118 119 # CHECK: 2120 print(s2.position)121 # CHECK: s2122 print(s2)123 124 assert s1 == s11125 assert s1 != s2126 127 128# CHECK-LABEL: TEST: testAffineAddExpr129@run130def testAffineAddExpr():131 with Context():132 d1 = AffineDimExpr.get(1)133 d2 = AffineDimExpr.get(2)134 d12 = AffineExpr.get_add(d1, d2)135 # CHECK: d1 + d2136 print(d12)137 138 d12op = d1 + d2139 # CHECK: d1 + d2140 print(d12op)141 142 d1cst_op = d1 + 2143 # CHECK: d1 + 2144 print(d1cst_op)145 146 d1cst_op2 = 2 + d1147 # CHECK: d1 + 2148 print(d1cst_op2)149 150 assert d12 == d12op151 assert d12.lhs == d1152 assert d12.rhs == d2153 154 155# CHECK-LABEL: TEST: testAffineMulExpr156@run157def testAffineMulExpr():158 with Context():159 d1 = AffineDimExpr.get(1)160 c2 = AffineConstantExpr.get(2)161 expr = AffineExpr.get_mul(d1, c2)162 # CHECK: d1 * 2163 print(expr)164 165 # CHECK: d1 * 2166 op = d1 * c2167 print(op)168 169 # CHECK: d1 * 2170 op_cst = d1 * 2171 print(op_cst)172 173 # CHECK: d1 * 2174 op_cst2 = 2 * d1175 print(op_cst2)176 177 assert expr == op178 assert expr == op_cst179 assert expr.lhs == d1180 assert expr.rhs == c2181 182 183# CHECK-LABEL: TEST: testAffineModExpr184@run185def testAffineModExpr():186 with Context():187 d1 = AffineDimExpr.get(1)188 c2 = AffineConstantExpr.get(2)189 expr = AffineExpr.get_mod(d1, c2)190 # CHECK: d1 mod 2191 print(expr)192 193 # CHECK: d1 mod 2194 op = d1 % c2195 print(op)196 197 # CHECK: d1 mod 2198 op_cst = d1 % 2199 print(op_cst)200 201 # CHECK: 2 mod d1202 print(2 % d1)203 204 assert expr == op205 assert expr == op_cst206 assert expr.lhs == d1207 assert expr.rhs == c2208 209 expr2 = AffineExpr.get_mod(c2, d1)210 expr3 = AffineExpr.get_mod(2, d1)211 expr4 = AffineExpr.get_mod(d1, 2)212 213 # CHECK: 2 mod d1214 print(expr2)215 # CHECK: 2 mod d1216 print(expr3)217 # CHECK: d1 mod 2218 print(expr4)219 220 assert expr2 == expr3221 assert expr4 == expr222 223 224# CHECK-LABEL: TEST: testAffineFloorDivExpr225@run226def testAffineFloorDivExpr():227 with Context():228 d1 = AffineDimExpr.get(1)229 c2 = AffineConstantExpr.get(2)230 expr = AffineExpr.get_floor_div(d1, c2)231 # CHECK: d1 floordiv 2232 print(expr)233 234 assert expr.lhs == d1235 assert expr.rhs == c2236 237 expr2 = AffineExpr.get_floor_div(c2, d1)238 expr3 = AffineExpr.get_floor_div(2, d1)239 expr4 = AffineExpr.get_floor_div(d1, 2)240 241 # CHECK: 2 floordiv d1242 print(expr2)243 # CHECK: 2 floordiv d1244 print(expr3)245 # CHECK: d1 floordiv 2246 print(expr4)247 248 assert expr2 == expr3249 assert expr4 == expr250 251 252# CHECK-LABEL: TEST: testAffineCeilDivExpr253@run254def testAffineCeilDivExpr():255 with Context():256 d1 = AffineDimExpr.get(1)257 c2 = AffineConstantExpr.get(2)258 expr = AffineExpr.get_ceil_div(d1, c2)259 # CHECK: d1 ceildiv 2260 print(expr)261 262 assert expr.lhs == d1263 assert expr.rhs == c2264 265 expr2 = AffineExpr.get_ceil_div(c2, d1)266 expr3 = AffineExpr.get_ceil_div(2, d1)267 expr4 = AffineExpr.get_ceil_div(d1, 2)268 269 # CHECK: 2 ceildiv d1270 print(expr2)271 # CHECK: 2 ceildiv d1272 print(expr3)273 # CHECK: d1 ceildiv 2274 print(expr4)275 276 assert expr2 == expr3277 assert expr4 == expr278 279 280# CHECK-LABEL: TEST: testAffineExprSub281@run282def testAffineExprSub():283 with Context():284 d1 = AffineDimExpr.get(1)285 d2 = AffineDimExpr.get(2)286 expr = d1 - d2287 # CHECK: d1 - d2288 print(expr)289 290 assert expr.lhs == d1291 rhs = AffineMulExpr(expr.rhs)292 # CHECK: d2293 print(rhs.lhs)294 # CHECK: -1295 print(rhs.rhs)296 297 # CHECK: d1 - 42298 print(d1 - 42)299 # CHECK: -d1 + 42300 print(42 - d1)301 302 c42 = AffineConstantExpr.get(42)303 assert d1 - 42 == d1 - c42304 assert 42 - d1 == c42 - d1305 306 307# CHECK-LABEL: TEST: testClassHierarchy308@run309def testClassHierarchy():310 with Context():311 d1 = AffineDimExpr.get(1)312 c2 = AffineConstantExpr.get(2)313 add = AffineAddExpr.get(d1, c2)314 mul = AffineMulExpr.get(d1, c2)315 mod = AffineModExpr.get(d1, c2)316 floor_div = AffineFloorDivExpr.get(d1, c2)317 ceil_div = AffineCeilDivExpr.get(d1, c2)318 319 # CHECK: False320 print(isinstance(d1, AffineBinaryExpr))321 # CHECK: False322 print(isinstance(c2, AffineBinaryExpr))323 # CHECK: True324 print(isinstance(add, AffineBinaryExpr))325 # CHECK: True326 print(isinstance(mul, AffineBinaryExpr))327 # CHECK: True328 print(isinstance(mod, AffineBinaryExpr))329 # CHECK: True330 print(isinstance(floor_div, AffineBinaryExpr))331 # CHECK: True332 print(isinstance(ceil_div, AffineBinaryExpr))333 334 try:335 AffineBinaryExpr(d1)336 except ValueError as e:337 # CHECK: Cannot cast affine expression to AffineBinaryExpr338 print(e)339 340 try:341 AffineBinaryExpr(c2)342 except ValueError as e:343 # CHECK: Cannot cast affine expression to AffineBinaryExpr344 print(e)345 346 347# CHECK-LABEL: TEST: testIsInstance348@run349def testIsInstance():350 with Context():351 d1 = AffineDimExpr.get(1)352 c2 = AffineConstantExpr.get(2)353 add = AffineAddExpr.get(d1, c2)354 mul = AffineMulExpr.get(d1, c2)355 356 # CHECK: True357 print(AffineDimExpr.isinstance(d1))358 # CHECK: False359 print(AffineConstantExpr.isinstance(d1))360 # CHECK: True361 print(AffineConstantExpr.isinstance(c2))362 # CHECK: False363 print(AffineMulExpr.isinstance(c2))364 # CHECK: True365 print(AffineAddExpr.isinstance(add))366 # CHECK: False367 print(AffineMulExpr.isinstance(add))368 # CHECK: True369 print(AffineMulExpr.isinstance(mul))370 # CHECK: False371 print(AffineAddExpr.isinstance(mul))372 373 374# CHECK-LABEL: TEST: testCompose375@run376def testCompose():377 with Context():378 # d0 + d2.379 expr = AffineAddExpr.get(AffineDimExpr.get(0), AffineDimExpr.get(2))380 381 # (d0, d1, d2)[s0, s1] -> (d0 + s1, d1 + s0, d0 + d1 + d2)382 map1 = AffineAddExpr.get(AffineDimExpr.get(0), AffineSymbolExpr.get(1))383 map2 = AffineAddExpr.get(AffineDimExpr.get(1), AffineSymbolExpr.get(0))384 map3 = AffineAddExpr.get(385 AffineAddExpr.get(AffineDimExpr.get(0), AffineDimExpr.get(1)),386 AffineDimExpr.get(2),387 )388 map = AffineMap.get(3, 2, [map1, map2, map3])389 390 # CHECK: d0 + s1 + d0 + d1 + d2391 print(expr.compose(map))392 393 394# CHECK-LABEL: TEST: testHash395@run396def testHash():397 with Context():398 d0 = AffineDimExpr.get(0)399 s1 = AffineSymbolExpr.get(1)400 assert hash(d0) == hash(AffineDimExpr.get(0))401 assert hash(d0 + s1) == hash(AffineAddExpr.get(d0, s1))402 403 dictionary = dict()404 dictionary[d0] = 0405 dictionary[s1] = 1406 assert d0 in dictionary407 assert s1 in dictionary408 409 410# CHECK-LABEL: TEST: testAffineExprShift411@run412def testAffineExprShift():413 with Context() as ctx:414 dims = [AffineExpr.get_dim(i) for i in range(4)]415 syms = [AffineExpr.get_symbol(i) for i in range(4)]416 417 assert (dims[2] + dims[3]) == (dims[0] + dims[1]).shift_dims(2, 2)418 assert (syms[2] + syms[3]) == (syms[0] + syms[1]).shift_symbols(2, 2, 0)419 420 421# CHECK-LABEL: TEST: testAffineExprSimplify422@run423def testAffineExprSimplify():424 with Context() as ctx:425 expr = AffineExpr.get_dim(0) + AffineExpr.get_symbol(0)426 assert expr == AffineExpr.simplify_affine_expr(expr, 1, 1)427