| # RUN: %PYTHON %s | FileCheck %s |
| |
| from mlir.ir import * |
| from mlir.dialects import func |
| from mlir.dialects import arith |
| from mlir.dialects import memref |
| from mlir.dialects import affine |
| import mlir.extras.types as T |
| |
| |
| def constructAndPrintInModule(f): |
| print("\nTEST:", f.__name__) |
| with Context(), Location.unknown(): |
| module = Module.create() |
| with InsertionPoint(module.body): |
| f() |
| print(module) |
| return f |
| |
| |
| # CHECK-LABEL: TEST: testAffineStoreOp |
| @constructAndPrintInModule |
| def testAffineStoreOp(): |
| f32 = F32Type.get() |
| index_type = IndexType.get() |
| memref_type_out = MemRefType.get([12, 12], f32) |
| |
| # CHECK: func.func @affine_store_test(%[[ARG0:.*]]: index) -> memref<12x12xf32> { |
| @func.FuncOp.from_py_func(index_type) |
| def affine_store_test(arg0): |
| # CHECK: %[[O_VAR:.*]] = memref.alloc() : memref<12x12xf32> |
| mem = memref.AllocOp(memref_type_out, [], []).result |
| |
| d0 = AffineDimExpr.get(0) |
| s0 = AffineSymbolExpr.get(0) |
| map = AffineMap.get(1, 1, [s0 * 3, d0 + s0 + 1]) |
| |
| # CHECK: %[[A1:.*]] = arith.constant 2.100000e+00 : f32 |
| a1 = arith.ConstantOp(f32, 2.1) |
| |
| # CHECK: affine.store %[[A1]], %alloc[symbol(%[[ARG0]]) * 3, %[[ARG0]] + symbol(%[[ARG0]]) + 1] : memref<12x12xf32> |
| affine.AffineStoreOp(a1, mem, indices=[arg0, arg0], map=map) |
| |
| return mem |
| |
| |
| # CHECK-LABEL: TEST: testAffineDelinearizeInfer |
| @constructAndPrintInModule |
| def testAffineDelinearizeInfer(): |
| # CHECK: %[[C1:.*]] = arith.constant 1 : index |
| c1 = arith.ConstantOp(T.index(), 1) |
| # CHECK: %{{.*}}:2 = affine.delinearize_index %[[C1:.*]] into (2, 3) : index, index |
| two_indices = affine.AffineDelinearizeIndexOp([T.index()] * 2, c1, [], [2, 3]) |
| |
| |
| # CHECK-LABEL: TEST: testAffineLoadOp |
| @constructAndPrintInModule |
| def testAffineLoadOp(): |
| f32 = F32Type.get() |
| index_type = IndexType.get() |
| memref_type_in = MemRefType.get([10, 10], f32) |
| |
| # CHECK: func.func @affine_load_test(%[[I_VAR:.*]]: memref<10x10xf32>, %[[ARG0:.*]]: index) -> f32 { |
| @func.FuncOp.from_py_func(memref_type_in, index_type) |
| def affine_load_test(I, arg0): |
| d0 = AffineDimExpr.get(0) |
| s0 = AffineSymbolExpr.get(0) |
| map = AffineMap.get(1, 1, [s0 * 3, d0 + s0 + 1]) |
| |
| # CHECK: {{.*}} = affine.load %[[I_VAR]][symbol(%[[ARG0]]) * 3, %[[ARG0]] + symbol(%[[ARG0]]) + 1] : memref<10x10xf32> |
| a1 = affine.AffineLoadOp(f32, I, indices=[arg0, arg0], map=map) |
| |
| return a1 |
| |
| |
| # CHECK-LABEL: TEST: testAffineForOp |
| @constructAndPrintInModule |
| def testAffineForOp(): |
| f32 = F32Type.get() |
| index_type = IndexType.get() |
| memref_type = MemRefType.get([1024], f32) |
| |
| # CHECK: #[[MAP0:.*]] = affine_map<(d0)[s0] -> (0, d0 + s0)> |
| # CHECK: #[[MAP1:.*]] = affine_map<(d0, d1) -> (d0 - 2, d1 * 32)> |
| # CHECK: func.func @affine_for_op_test(%[[BUFFER:.*]]: memref<1024xf32>) { |
| @func.FuncOp.from_py_func(memref_type) |
| def affine_for_op_test(buffer): |
| # CHECK: %[[C1:.*]] = arith.constant 1 : index |
| c1 = arith.ConstantOp(index_type, 1) |
| # CHECK: %[[C2:.*]] = arith.constant 2 : index |
| c2 = arith.ConstantOp(index_type, 2) |
| # CHECK: %[[C3:.*]] = arith.constant 3 : index |
| c3 = arith.ConstantOp(index_type, 3) |
| # CHECK: %[[C9:.*]] = arith.constant 9 : index |
| c9 = arith.ConstantOp(index_type, 9) |
| # CHECK: %[[AC0:.*]] = arith.constant 0.000000e+00 : f32 |
| ac0 = AffineConstantExpr.get(0) |
| |
| d0 = AffineDimExpr.get(0) |
| d1 = AffineDimExpr.get(1) |
| s0 = AffineSymbolExpr.get(0) |
| lb = AffineMap.get(1, 1, [ac0, d0 + s0]) |
| ub = AffineMap.get(2, 0, [d0 - 2, 32 * d1]) |
| sum_0 = arith.ConstantOp(f32, 0.0) |
| |
| # CHECK: %0 = affine.for %[[INDVAR:.*]] = max #[[MAP0]](%[[C2]])[%[[C3]]] to min #[[MAP1]](%[[C9]], %[[C1]]) step 2 iter_args(%[[SUM0:.*]] = %[[AC0]]) -> (f32) { |
| sum = affine.AffineForOp( |
| lb, |
| ub, |
| 2, |
| iter_args=[sum_0], |
| lower_bound_operands=[c2, c3], |
| upper_bound_operands=[c9, c1], |
| ) |
| |
| with InsertionPoint(sum.body): |
| # CHECK: %[[TMP:.*]] = memref.load %[[BUFFER]][%[[INDVAR]]] : memref<1024xf32> |
| tmp = memref.LoadOp(buffer, [sum.induction_variable]) |
| sum_next = arith.AddFOp(sum.inner_iter_args[0], tmp) |
| affine.AffineYieldOp([sum_next]) |
| |
| |
| # CHECK-LABEL: TEST: testAffineForOpErrors |
| @constructAndPrintInModule |
| def testAffineForOpErrors(): |
| c1 = arith.ConstantOp(T.index(), 1) |
| c2 = arith.ConstantOp(T.index(), 2) |
| c3 = arith.ConstantOp(T.index(), 3) |
| d0 = AffineDimExpr.get(0) |
| |
| try: |
| affine.AffineForOp( |
| c1, |
| c2, |
| 1, |
| lower_bound_operands=[c3], |
| upper_bound_operands=[], |
| ) |
| except ValueError as e: |
| assert ( |
| e.args[0] |
| == "Either a concrete lower bound or an AffineMap in combination with lower bound operands, but not both, is supported." |
| ) |
| |
| try: |
| affine.AffineForOp( |
| AffineMap.get_constant(1), |
| c2, |
| 1, |
| lower_bound_operands=[c3, c3], |
| upper_bound_operands=[], |
| ) |
| except ValueError as e: |
| assert ( |
| e.args[0] |
| == "Wrong number of lower bound operands passed to AffineForOp; Expected 0, got 2." |
| ) |
| |
| try: |
| two_indices = affine.AffineDelinearizeIndexOp([T.index()] * 2, c1, [], [1, 1]) |
| affine.AffineForOp( |
| two_indices, |
| c2, |
| 1, |
| lower_bound_operands=[], |
| upper_bound_operands=[], |
| ) |
| except ValueError as e: |
| assert e.args[0] == "Only a single concrete value is supported for lower bound." |
| |
| try: |
| affine.AffineForOp( |
| 1.0, |
| c2, |
| 1, |
| lower_bound_operands=[], |
| upper_bound_operands=[], |
| ) |
| except ValueError as e: |
| assert e.args[0] == "lower bound must be int | ResultValueT | AffineMap." |
| |
| |
| @constructAndPrintInModule |
| def testForSugar(): |
| memref_t = T.memref(10, T.index()) |
| range = affine.for_ |
| |
| # CHECK: #[[$ATTR_2:.+]] = affine_map<(d0) -> (d0)> |
| |
| # CHECK-LABEL: func.func @range_loop_1( |
| # CHECK-SAME: %[[VAL_0:.*]]: index, %[[VAL_1:.*]]: index, %[[VAL_2:.*]]: memref<10xindex>) { |
| # CHECK: affine.for %[[VAL_3:.*]] = #[[$ATTR_2]](%[[VAL_0]]) to #[[$ATTR_2]](%[[VAL_1]]) { |
| # CHECK: %[[VAL_4:.*]] = arith.addi %[[VAL_3]], %[[VAL_3]] : index |
| # CHECK: memref.store %[[VAL_4]], %[[VAL_2]]{{\[}}%[[VAL_3]]] : memref<10xindex> |
| # CHECK: } |
| # CHECK: return |
| # CHECK: } |
| @func.FuncOp.from_py_func(T.index(), T.index(), memref_t) |
| def range_loop_1(lb, ub, memref_v): |
| for i in range(lb, ub, step=1): |
| add = arith.addi(i, i) |
| memref.store(add, memref_v, [i]) |
| |
| affine.yield_([]) |
| |
| # CHECK-LABEL: func.func @range_loop_2( |
| # CHECK-SAME: %[[VAL_0:.*]]: index, %[[VAL_1:.*]]: index, %[[VAL_2:.*]]: memref<10xindex>) { |
| # CHECK: affine.for %[[VAL_3:.*]] = #[[$ATTR_2]](%[[VAL_0]]) to 10 { |
| # CHECK: %[[VAL_4:.*]] = arith.addi %[[VAL_3]], %[[VAL_3]] : index |
| # CHECK: memref.store %[[VAL_4]], %[[VAL_2]]{{\[}}%[[VAL_3]]] : memref<10xindex> |
| # CHECK: } |
| # CHECK: return |
| # CHECK: } |
| @func.FuncOp.from_py_func(T.index(), T.index(), memref_t) |
| def range_loop_2(lb, ub, memref_v): |
| for i in range(lb, 10, step=1): |
| add = arith.addi(i, i) |
| memref.store(add, memref_v, [i]) |
| affine.yield_([]) |
| |
| # CHECK-LABEL: func.func @range_loop_3( |
| # CHECK-SAME: %[[VAL_0:.*]]: index, %[[VAL_1:.*]]: index, %[[VAL_2:.*]]: memref<10xindex>) { |
| # CHECK: affine.for %[[VAL_3:.*]] = 0 to #[[$ATTR_2]](%[[VAL_1]]) { |
| # CHECK: %[[VAL_4:.*]] = arith.addi %[[VAL_3]], %[[VAL_3]] : index |
| # CHECK: memref.store %[[VAL_4]], %[[VAL_2]]{{\[}}%[[VAL_3]]] : memref<10xindex> |
| # CHECK: } |
| # CHECK: return |
| # CHECK: } |
| @func.FuncOp.from_py_func(T.index(), T.index(), memref_t) |
| def range_loop_3(lb, ub, memref_v): |
| for i in range(0, ub, step=1): |
| add = arith.addi(i, i) |
| memref.store(add, memref_v, [i]) |
| affine.yield_([]) |
| |
| # CHECK-LABEL: func.func @range_loop_4( |
| # CHECK-SAME: %[[VAL_0:.*]]: index, %[[VAL_1:.*]]: index, %[[VAL_2:.*]]: memref<10xindex>) { |
| # CHECK: affine.for %[[VAL_3:.*]] = 0 to 10 { |
| # CHECK: %[[VAL_4:.*]] = arith.addi %[[VAL_3]], %[[VAL_3]] : index |
| # CHECK: memref.store %[[VAL_4]], %[[VAL_2]]{{\[}}%[[VAL_3]]] : memref<10xindex> |
| # CHECK: } |
| # CHECK: return |
| # CHECK: } |
| @func.FuncOp.from_py_func(T.index(), T.index(), memref_t) |
| def range_loop_4(lb, ub, memref_v): |
| for i in range(0, 10, step=1): |
| add = arith.addi(i, i) |
| memref.store(add, memref_v, [i]) |
| affine.yield_([]) |
| |
| # CHECK-LABEL: func.func @range_loop_8( |
| # CHECK-SAME: %[[VAL_0:.*]]: index, %[[VAL_1:.*]]: index, %[[VAL_2:.*]]: memref<10xindex>) { |
| # CHECK: %[[VAL_3:.*]] = affine.for %[[VAL_4:.*]] = 0 to 10 iter_args(%[[VAL_5:.*]] = %[[VAL_2]]) -> (memref<10xindex>) { |
| # CHECK: %[[VAL_6:.*]] = arith.addi %[[VAL_4]], %[[VAL_4]] : index |
| # CHECK: memref.store %[[VAL_6]], %[[VAL_5]]{{\[}}%[[VAL_4]]] : memref<10xindex> |
| # CHECK: affine.yield %[[VAL_5]] : memref<10xindex> |
| # CHECK: } |
| # CHECK: return |
| # CHECK: } |
| @func.FuncOp.from_py_func(T.index(), T.index(), memref_t) |
| def range_loop_8(lb, ub, memref_v): |
| for i, it in range(0, 10, iter_args=[memref_v]): |
| add = arith.addi(i, i) |
| memref.store(add, it, [i]) |
| affine.yield_([it]) |
| |
| |
| # CHECK-LABEL: TEST: testAffineIfWithoutElse |
| @constructAndPrintInModule |
| def testAffineIfWithoutElse(): |
| index = IndexType.get() |
| i32 = IntegerType.get_signless(32) |
| d0 = AffineDimExpr.get(0) |
| |
| # CHECK: #[[$SET0:.*]] = affine_set<(d0) : (d0 - 5 >= 0)> |
| cond = IntegerSet.get(1, 0, [d0 - 5], [False]) |
| |
| # CHECK-LABEL: func.func @simple_affine_if( |
| # CHECK-SAME: %[[VAL_0:.*]]: index) { |
| # CHECK: affine.if #[[$SET0]](%[[VAL_0]]) { |
| # CHECK: %[[VAL_1:.*]] = arith.constant 1 : i32 |
| # CHECK: %[[VAL_2:.*]] = arith.addi %[[VAL_1]], %[[VAL_1]] : i32 |
| # CHECK: } |
| # CHECK: return |
| # CHECK: } |
| @func.FuncOp.from_py_func(index) |
| def simple_affine_if(cond_operands): |
| if_op = affine.AffineIfOp(cond, cond_operands=[cond_operands]) |
| with InsertionPoint(if_op.then_block): |
| one = arith.ConstantOp(i32, 1) |
| add = arith.AddIOp(one, one) |
| affine.AffineYieldOp([]) |
| return |
| |
| |
| # CHECK-LABEL: TEST: testAffineIfWithElse |
| @constructAndPrintInModule |
| def testAffineIfWithElse(): |
| index = IndexType.get() |
| i32 = IntegerType.get_signless(32) |
| d0 = AffineDimExpr.get(0) |
| |
| # CHECK: #[[$SET0:.*]] = affine_set<(d0) : (d0 - 5 >= 0)> |
| cond = IntegerSet.get(1, 0, [d0 - 5], [False]) |
| |
| # CHECK-LABEL: func.func @simple_affine_if_else( |
| # CHECK-SAME: %[[VAL_0:.*]]: index) { |
| # CHECK: %[[VAL_IF:.*]]:2 = affine.if #[[$SET0]](%[[VAL_0]]) -> (i32, i32) { |
| # CHECK: %[[VAL_XT:.*]] = arith.constant 0 : i32 |
| # CHECK: %[[VAL_YT:.*]] = arith.constant 1 : i32 |
| # CHECK: affine.yield %[[VAL_XT]], %[[VAL_YT]] : i32, i32 |
| # CHECK: } else { |
| # CHECK: %[[VAL_XF:.*]] = arith.constant 2 : i32 |
| # CHECK: %[[VAL_YF:.*]] = arith.constant 3 : i32 |
| # CHECK: affine.yield %[[VAL_XF]], %[[VAL_YF]] : i32, i32 |
| # CHECK: } |
| # CHECK: %[[VAL_ADD:.*]] = arith.addi %[[VAL_IF]]#0, %[[VAL_IF]]#1 : i32 |
| # CHECK: return |
| # CHECK: } |
| |
| @func.FuncOp.from_py_func(index) |
| def simple_affine_if_else(cond_operands): |
| if_op = affine.AffineIfOp( |
| cond, [i32, i32], cond_operands=[cond_operands], has_else=True |
| ) |
| with InsertionPoint(if_op.then_block): |
| x_true = arith.ConstantOp(i32, 0) |
| y_true = arith.ConstantOp(i32, 1) |
| affine.AffineYieldOp([x_true, y_true]) |
| with InsertionPoint(if_op.else_block): |
| x_false = arith.ConstantOp(i32, 2) |
| y_false = arith.ConstantOp(i32, 3) |
| affine.AffineYieldOp([x_false, y_false]) |
| add = arith.AddIOp(if_op.results[0], if_op.results[1]) |
| return |
