| # RUN: %PYTHON %s | FileCheck %s |
| |
| import gc |
| from mlir.ir import * |
| |
| |
| def run(f): |
| print("\nTEST:", f.__name__) |
| f() |
| gc.collect() |
| assert Context._get_live_count() == 0 |
| return f |
| |
| |
| # CHECK-LABEL: TEST: testAffineMapCapsule |
| @run |
| def testAffineMapCapsule(): |
| with Context() as ctx: |
| am1 = AffineMap.get_empty(ctx) |
| # CHECK: mlir.ir.AffineMap._CAPIPtr |
| affine_map_capsule = am1._CAPIPtr |
| print(affine_map_capsule) |
| am2 = AffineMap._CAPICreate(affine_map_capsule) |
| assert am2 == am1 |
| assert am2.context is ctx |
| |
| |
| # CHECK-LABEL: TEST: testAffineMapGet |
| @run |
| def testAffineMapGet(): |
| with Context() as ctx: |
| d0 = AffineDimExpr.get(0) |
| d1 = AffineDimExpr.get(1) |
| c2 = AffineConstantExpr.get(2) |
| |
| # CHECK: (d0, d1)[s0, s1, s2] -> () |
| map0 = AffineMap.get(2, 3, []) |
| print(map0) |
| |
| # CHECK: (d0, d1)[s0, s1, s2] -> (d1, 2) |
| map1 = AffineMap.get(2, 3, [d1, c2]) |
| print(map1) |
| |
| # CHECK: () -> (2) |
| map2 = AffineMap.get(0, 0, [c2]) |
| print(map2) |
| |
| # CHECK: (d0, d1) -> (d0, d1) |
| map3 = AffineMap.get(2, 0, [d0, d1]) |
| print(map3) |
| |
| # CHECK: (d0, d1) -> (d1) |
| map4 = AffineMap.get(2, 0, [d1]) |
| print(map4) |
| |
| # CHECK: (d0, d1, d2) -> (d2, d0, d1) |
| map5 = AffineMap.get_permutation([2, 0, 1]) |
| print(map5) |
| |
| assert map1 == AffineMap.get(2, 3, [d1, c2]) |
| assert AffineMap.get(0, 0, []) == AffineMap.get_empty() |
| assert map2 == AffineMap.get_constant(2) |
| assert map3 == AffineMap.get_identity(2) |
| assert map4 == AffineMap.get_minor_identity(2, 1) |
| |
| try: |
| AffineMap.get(1, 1, [1]) |
| except RuntimeError as e: |
| # CHECK: Invalid expression when attempting to create an AffineMap |
| print(e) |
| |
| try: |
| AffineMap.get(1, 1, [None]) |
| except RuntimeError as e: |
| # CHECK: Invalid expression (None?) when attempting to create an AffineMap |
| print(e) |
| |
| try: |
| AffineMap.get_permutation([1, 0, 1]) |
| except RuntimeError as e: |
| # CHECK: Invalid permutation when attempting to create an AffineMap |
| print(e) |
| |
| try: |
| map3.get_submap([42]) |
| except ValueError as e: |
| # CHECK: result position out of bounds |
| print(e) |
| |
| try: |
| map3.get_minor_submap(42) |
| except ValueError as e: |
| # CHECK: number of results out of bounds |
| print(e) |
| |
| try: |
| map3.get_major_submap(42) |
| except ValueError as e: |
| # CHECK: number of results out of bounds |
| print(e) |
| |
| |
| # CHECK-LABEL: TEST: testAffineMapDerive |
| @run |
| def testAffineMapDerive(): |
| with Context() as ctx: |
| map5 = AffineMap.get_identity(5) |
| |
| # CHECK: (d0, d1, d2, d3, d4) -> (d1, d2, d3) |
| map123 = map5.get_submap([1, 2, 3]) |
| print(map123) |
| |
| # CHECK: (d0, d1, d2, d3, d4) -> (d0, d1) |
| map01 = map5.get_major_submap(2) |
| print(map01) |
| |
| # CHECK: (d0, d1, d2, d3, d4) -> (d3, d4) |
| map34 = map5.get_minor_submap(2) |
| print(map34) |
| |
| |
| # CHECK-LABEL: TEST: testAffineMapProperties |
| @run |
| def testAffineMapProperties(): |
| with Context(): |
| d0 = AffineDimExpr.get(0) |
| d1 = AffineDimExpr.get(1) |
| d2 = AffineDimExpr.get(2) |
| map1 = AffineMap.get(3, 0, [d2, d0]) |
| map2 = AffineMap.get(3, 0, [d2, d0, d1]) |
| map3 = AffineMap.get(3, 1, [d2, d0, d1]) |
| # CHECK: False |
| print(map1.is_permutation) |
| # CHECK: True |
| print(map1.is_projected_permutation) |
| # CHECK: True |
| print(map2.is_permutation) |
| # CHECK: True |
| print(map2.is_projected_permutation) |
| # CHECK: False |
| print(map3.is_permutation) |
| # CHECK: False |
| print(map3.is_projected_permutation) |
| |
| |
| # CHECK-LABEL: TEST: testAffineMapExprs |
| @run |
| def testAffineMapExprs(): |
| with Context(): |
| d0 = AffineDimExpr.get(0) |
| d1 = AffineDimExpr.get(1) |
| d2 = AffineDimExpr.get(2) |
| map3 = AffineMap.get(3, 1, [d2, d0, d1]) |
| |
| # CHECK: 3 |
| print(map3.n_dims) |
| # CHECK: 4 |
| print(map3.n_inputs) |
| # CHECK: 1 |
| print(map3.n_symbols) |
| assert map3.n_inputs == map3.n_dims + map3.n_symbols |
| |
| # CHECK: 3 |
| print(len(map3.results)) |
| for expr in map3.results: |
| # CHECK: d2 |
| # CHECK: d0 |
| # CHECK: d1 |
| print(expr) |
| for expr in map3.results[-1:-4:-1]: |
| # CHECK: d1 |
| # CHECK: d0 |
| # CHECK: d2 |
| print(expr) |
| assert list(map3.results) == [d2, d0, d1] |
| |
| |
| # CHECK-LABEL: TEST: testCompressUnusedSymbols |
| @run |
| def testCompressUnusedSymbols(): |
| with Context() as ctx: |
| d0, d1, d2 = (AffineDimExpr.get(0), AffineDimExpr.get(1), |
| AffineDimExpr.get(2)) |
| s0, s1, s2 = (AffineSymbolExpr.get(0), AffineSymbolExpr.get(1), |
| AffineSymbolExpr.get(2)) |
| maps = [ |
| AffineMap.get(3, 3, [d2, d0, d1]), |
| AffineMap.get(3, 3, [d2, d0 + s2, d1]), |
| AffineMap.get(3, 3, [d1, d2, d0]) |
| ] |
| |
| compressed_maps = AffineMap.compress_unused_symbols(maps, ctx) |
| |
| # CHECK: AffineMap((d0, d1, d2)[s0, s1, s2] -> (d2, d0, d1)) |
| # CHECK-SAME: AffineMap((d0, d1, d2)[s0, s1, s2] -> (d2, d0 + s2, d1)) |
| # CHECK-SAME: AffineMap((d0, d1, d2)[s0, s1, s2] -> (d1, d2, d0)) |
| print(maps) |
| |
| # CHECK: AffineMap((d0, d1, d2)[s0] -> (d2, d0, d1)) |
| # CHECK-SAME: AffineMap((d0, d1, d2)[s0] -> (d2, d0 + s0, d1)) |
| # CHECK-SAME: AffineMap((d0, d1, d2)[s0] -> (d1, d2, d0)) |
| print(compressed_maps) |
| |
| |
| # CHECK-LABEL: TEST: testReplace |
| @run |
| def testReplace(): |
| with Context() as ctx: |
| d0, d1, d2 = (AffineDimExpr.get(0), AffineDimExpr.get(1), |
| AffineDimExpr.get(2)) |
| s0, s1, s2 = (AffineSymbolExpr.get(0), AffineSymbolExpr.get(1), |
| AffineSymbolExpr.get(2)) |
| map1 = AffineMap.get(3, 3, [d2, d0 + s1 + s2, d1 + s0]) |
| |
| replace0 = map1.replace(s0, AffineConstantExpr.get(42), 3, 3) |
| replace1 = map1.replace(s1, AffineConstantExpr.get(42), 3, 3) |
| replace3 = map1.replace(s2, AffineConstantExpr.get(42), 3, 2) |
| |
| # CHECK: (d0, d1, d2)[s0, s1, s2] -> (d2, d0 + s1 + s2, d1 + 42) |
| print(replace0) |
| |
| # CHECK: (d0, d1, d2)[s0, s1, s2] -> (d2, d0 + s2 + 42, d1 + s0) |
| print(replace1) |
| |
| # CHECK: (d0, d1, d2)[s0, s1] -> (d2, d0 + s1 + 42, d1 + s0) |
| print(replace3) |
| |
| |
| # CHECK-LABEL: TEST: testHash |
| @run |
| def testHash(): |
| with Context(): |
| d0, d1 = AffineDimExpr.get(0), AffineDimExpr.get(1) |
| m1 = AffineMap.get(2, 0, [d0, d1]) |
| m2 = AffineMap.get(2, 0, [d1, d0]) |
| assert hash(m1) == hash(AffineMap.get(2, 0, [d0, d1])) |
| |
| dictionary = dict() |
| dictionary[m1] = 1 |
| dictionary[m2] = 2 |
| assert m1 in dictionary |