| # 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: %[[C0:.*]] = arith.constant 0 : index |
| c0 = arith.ConstantOp(T.index(), 0) |
| # CHECK: %[[C1:.*]] = arith.constant 1 : index |
| c1 = arith.ConstantOp(T.index(), 1) |
| # CHECK: %{{.*}}:2 = affine.delinearize_index %[[C1:.*]] into (%[[C1:.*]], %[[C0:.*]]) : index, index |
| two_indices = affine.AffineDelinearizeIndexOp(c1, [c1, c0]) |
| |
| |
| # 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(c1, [c1, c1]) |
| 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]) |
