mlir/test/python/dialects/affine.py

# 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


@constructAndPrintInModule
def test_double_AtomicRMWKindAttr_registration():
    from mlir.dialects import _affine_enum_gen


# CHECK-LABEL: TEST: testAffineIfOpInsertionPoint
@constructAndPrintInModule
def testAffineIfOpInsertionPoint():
    index = IndexType.get()
    i32 = IntegerType.get_signless(32)
    d0 = AffineDimExpr.get(0)
    cond = IntegerSet.get(1, 0, [d0 - 5], [False])

    # CHECK: func.func @affine_if_insertion_point_test
    @func.FuncOp.from_py_func(index)
    def affine_if_insertion_point_test(cond_operands):
        first_const = arith.ConstantOp(i32, 42)
        ip_before = InsertionPoint(first_const)
        second_const = arith.ConstantOp(i32, 100)

        # Now create AffineIfOp with ip pointing BEFORE first_const
        # If ip works correctly, the if_op should appear BEFORE first_const
        if_op = affine.AffineIfOp(cond, cond_operands=[cond_operands], ip=ip_before)
        with InsertionPoint(if_op.then_block):
            affine.AffineYieldOp([])

        # CHECK: affine.if
        # CHECK-NEXT: }
        # CHECK-NEXT: %c42_i32 = arith.constant 42 : i32
        # CHECK-NEXT: %c100_i32 = arith.constant 100 : i32
        return