# RUN: %PYTHON %s | FileCheck %s

import functools
from typing import Callable

from mlir.ir import *
from mlir.dialects import transform
from mlir.dialects import pdl
from mlir.dialects.transform import structured
from mlir.dialects.transform import pdl as transform_pdl


def run(f):
    with Context(), Location.unknown():
        module = Module.create()
        with InsertionPoint(module.body):
            print("\nTEST:", f.__name__)
            f()
        module.operation.verify()
        print(module)
    return f


def create_sequence(func: Callable) -> Callable:
    @functools.wraps(func)
    def decorated() -> None:
        sequence = transform.SequenceOp(
            transform.FailurePropagationMode.Propagate,
            [],
            transform.AnyOpType.get(),
        )
        with InsertionPoint(sequence.body):
            func(sequence.bodyTarget)
            transform.YieldOp()

    return decorated


@run
@create_sequence
def testBufferizeToAllocationOpCompact(target):
    structured.BufferizeToAllocationOp(target)
    # CHECK-LABEL: TEST: testBufferizeToAllocationOpCompact
    # CHECK: transform.sequence
    # CHECK: transform.structured.bufferize_to_allocation


@run
@create_sequence
def testBufferizeToAllocationOpArgs(target):
    structured.BufferizeToAllocationOp(
        target,
        memory_space=3,
        memcpy_op="memref.copy",
        alloc_op="memref.alloca",
        bufferize_destination_only=True,
    )
    # CHECK-LABEL: TEST: testBufferizeToAllocationOpArgs
    # CHECK: transform.sequence
    # CHECK: transform.structured.bufferize_to_allocation
    # CHECK-SAME: alloc_op = "memref.alloca"
    # CHECK-SAME: bufferize_destination_only
    # CHECK-SAME: memcpy_op = "memref.copy"
    # CHECK-SAME: memory_space = 3


@run
@create_sequence
def testDecompose(target):
    structured.DecomposeOp(target)
    # CHECK-LABEL: TEST: testDecompose
    # CHECK: transform.sequence
    # CHECK: transform.structured.decompose


@run
@create_sequence
def testFuseIntoContainingOpTypes(target):
    fused = structured.MatchOp.match_op_names(target, ["test.dummy"])
    containing = structured.MatchOp.match_op_names(target, ["test.dummy"])
    structured.FuseIntoContainingOp(
        transform.OperationType.get("test.dummy"),
        transform.OperationType.get("test.dummy"),
        fused,
        containing,
    )
    # CHECK-LABEL: TEST: testFuseIntoContainingOpTypes
    # CHECK: = transform.structured.fuse_into_containing_op
    # CHECK-SAME: (!transform.any_op, !transform.any_op) -> (!transform.op<"test.dummy">, !transform.op<"test.dummy">)


@run
@create_sequence
def testFuseIntoContainingOpCompact(target):
    fused = structured.MatchOp.match_op_names(target, ["test.dummy"])
    containing = structured.MatchOp.match_op_names(target, ["test.dummy"])
    structured.FuseIntoContainingOp(fused, containing)
    # CHECK-LABEL: TEST: testFuseIntoContainingOpCompact
    # CHECK: = transform.structured.fuse_into_containing_op
    # CHECK-SAME: (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)


@run
@create_sequence
def testGeneralize(target):
    structured.GeneralizeOp(target)
    # CHECK-LABEL: TEST: testGeneralize
    # CHECK: transform.sequence
    # CHECK: transform.structured.generalize


@run
@create_sequence
def testInterchange(target):
    structured.InterchangeOp(target, iterator_interchange=[1, 0])
    # CHECK-LABEL: TEST: testInterchange
    # CHECK: transform.sequence
    # CHECK: transform.structured.interchange
    # CHECK: iterator_interchange = [1, 0]


@run
@create_sequence
def testMapCopyToThreadsOpCompact(target):
    structured.MapCopyToThreadsOp(
        target, total_num_threads=32, desired_bit_alignment=128
    )
    # CHECK-LABEL: TEST: testMapCopyToThreadsOpCompact
    # CHECK: = transform.structured.gpu.map_copy_to_threads
    # CHECK-SAME: total_num_threads = 32
    # CHECK-SAME: desired_bit_alignment = 128
    # CHECK-SAME:  (!transform.any_op) -> (!transform.any_op, !transform.any_op)


@run
@create_sequence
def testMapCopyToThreadsOpTypes(target):
    structured.MapCopyToThreadsOp(
        transform.OperationType.get("test.opA"),
        transform.OperationType.get("test.opB"),
        target,
        total_num_threads=32,
        desired_bit_alignment=128,
    )
    # CHECK-LABEL: TEST: testMapCopyToThreadsOpTypes
    # CHECK: = transform.structured.gpu.map_copy_to_threads
    # CHECK-SAME: total_num_threads = 32
    # CHECK-SAME: desired_bit_alignment = 128
    # CHECK-SAME:  (!transform.any_op) -> (!transform.op<"test.opA">, !transform.op<"test.opB">)


@run
@create_sequence
def testMatchOpNamesString(target):
    structured.MatchOp.match_op_names(target, "test.dummy")
    # CHECK-LABEL: TEST: testMatchOpNamesString
    # CHECK: transform.structured.match ops
    # CHECK-SAME: ["test.dummy"]
    # CHECK-SAME: (!transform.any_op) -> !transform.any_op


@run
@create_sequence
def testMatchOpNamesList(target):
    structured.MatchOp.match_op_names(target, ["test.dummy"])
    # CHECK-LABEL: TEST: testMatchOpNamesList
    # CHECK: transform.structured.match ops
    # CHECK-SAME: ["test.dummy"]
    # CHECK-SAME: (!transform.any_op) -> !transform.any_op


@run
@create_sequence
def testVectorizeNoArgs(target):
    structured.VectorizeOp(target)
    # CHECK-LABEL: TEST: testVectorizeNoArgs
    # CHECK: transform.sequence
    # CHECK: transform.structured.vectorize
    # CHECK-NOT:     vector_sizes


@run
@create_sequence
def testVectorizeStatic(target):
    structured.VectorizeOp(target, [16, 4])
    # CHECK-LABEL: TEST: testVectorizeStatic
    # CHECK: transform.sequence
    # CHECK: transform.structured.vectorize
    # CHECK-SAME:     vector_sizes [16, 4]


@run
@create_sequence
def testVectorizeArray(target):
    sizes = Attribute.parse("[16, 4]")
    structured.VectorizeOp(target, sizes)
    # CHECK-LABEL: TEST: testVectorizeArray
    # CHECK: transform.sequence
    # CHECK: transform.structured.vectorize
    # CHECK-SAME:     vector_sizes [16, 4]


@run
@create_sequence
def testVectorizeMixed(target):
    sz1 = structured.MatchOp.match_op_names(target, ["arith.constant"])
    sz2 = Attribute.parse("4")
    structured.VectorizeOp(target, [sz1, sz2])
    # CHECK-LABEL: TEST: testVectorizeMixed
    # CHECK: transform.sequence
    # CHECK: %[[V0:.*]] = transform.structured.match
    # CHECK: transform.structured.vectorize
    # CHECK-SAME:     vector_sizes [%[[V0]], 4]


@run
@create_sequence
def testVectorizeEmpty(target):
    structured.VectorizeOp(target, [])
    # CHECK-LABEL: TEST: testVectorizeEmpty
    # CHECK: transform.sequence
    # CHECK: transform.structured.vectorize
    # CHECK-NOT:     vector_sizes


@run
@create_sequence
def testVectorizeScalable(target):
    sz1 = structured.MatchOp.match_op_names(target, ["arith.constant"])
    sz2 = Attribute.parse("4")
    structured.VectorizeOp(target, [16, [sz1], [sz2], [8]])
    # CHECK-LABEL: TEST: testVectorizeScalable
    # CHECK: transform.sequence
    # CHECK-DAG: %[[V0:.*]] = transform.structured.match
    # CHECK-DAG: transform.structured.vectorize
    # CHECK-SAME:     vector_sizes [16, [%[[V0]]], [4], [8]]


@run
@create_sequence
def testVectorizeArgs(target):
    structured.VectorizeOp(target, [16, 4], vectorize_nd_extract=True)
    # CHECK-LABEL: TEST: testVectorizeArgs
    # CHECK: transform.sequence
    # CHECK: transform.structured.vectorize
    # CHECK-SAME: vectorize_nd_extract


@run
@create_sequence
def testMatchOpNamesTyped(target):
    structured.MatchOp.match_op_names(
        transform.OperationType.get("test.dummy"),
        target,
        ["test.dummy"],
    )
    # CHECK-LABEL: TEST: testMatchOpNamesTyped
    # CHECK: transform.structured.match ops
    # CHECK-SAME: ["test.dummy"]
    # CHECK-SAME: (!transform.any_op) -> !transform.op<"test.dummy">


@run
@create_sequence
def testMultitileSizesCompact(target):
    structured.MultiTileSizesOp(
        transform.AnyOpType.get(), target, dimension=1, target_size=42
    )
    # CHECK-LABEL: TEST: testMultitileSizes
    # CHECK: transform.sequence
    # CHECK-NOT: divisor
    # CHECK: transform.structured.multitile_sizes
    # CHECK-NOT: divisor
    # CHECK-DAG: dimension = 1
    # CHECK-NOT: divisor
    # CHECK-DAG: target_size = 42
    # CHECK-NOT: divisor


@run
@create_sequence
def testMultitileSizesAllArgs(target):
    structured.MultiTileSizesOp(
        transform.AnyOpType.get(),
        target,
        dimension=1,
        target_size=42,
        divisor=2,
    )
    # CHECK-LABEL: TEST: testMultitileSizes
    # CHECK: transform.sequence
    # CHECK: transform.structured.multitile_sizes
    # CHECK-DAG: dimension = 1
    # CHECK-DAG: divisor = 2
    # CHECK-DAG: target_size = 42


@run
@create_sequence
def testPadOpNoArgs(target):
    structured.PadOp(target)
    # CHECK-LABEL: TEST: testPadOpNoArgs
    # CHECK: transform.sequence
    # CHECK: transform.structured.pad
    # CHECK-NOT: copy_back_op
    # CHECK-NOT: pack_paddings
    # CHECK-NOT: pad_to_multiple_of
    # CHECK-NOT: padding_dimensions
    # CHECK-NOT: padding_values
    # CHECK-NOT: transpose_paddings


@run
@create_sequence
def testPadOpArgs(target):
    structured.PadOp(
        target,
        padding_values=[FloatAttr.get_f32(42.0), StringAttr.get("0")],
        padding_dimensions=Attribute.parse("[1]"),
        pad_to_multiple_of=[128],
        pack_paddings=[0],
        transpose_paddings=[[1, Attribute.parse("0")], Attribute.parse("[0, 1]")],
        copy_back_op="linalg.copy",
    )
    # CHECK-LABEL: TEST: testPadOpArgs
    # CHECK: transform.sequence
    # CHECK: transform.structured.pad
    # CHECK-DAG: copy_back_op = "linalg.copy"
    # CHECK-DAG: pack_paddings = [0]
    # CHECK-DAG: pad_to_multiple_of = [128]
    # CHECK-DAG: padding_dimensions = [1]
    # CHECK-DAG: padding_values = [4.200000e+01 : f32, "0"]
    # CHECK-DAG: transpose_paddings = {{\[}}[1, 0], [0, 1]]


@run
@create_sequence
def testScalarize(target):
    structured.ScalarizeOp(target)
    # CHECK-LABEL: TEST: testScalarize
    # CHECK: transform.structured.scalarize


@run
@create_sequence
def testSplit(target):
    split = structured.SplitOp(target, dimension=1, split_point=42)
    structured.SplitOp(split.results[0], dimension=3, split_point=split.results[1])
    # CHECK-LABEL: TEST: testSplit
    # CHECK: %[[F:.+]], %[[S:.+]] = transform.structured.split %{{.*}} after 42 {dimension = 1
    # CHECK: transform.structured.split %[[F]] after %[[S]] {dimension = 3


@run
@create_sequence
def testTileCompact(target):
    structured.TileUsingForOp(target, sizes=[4, 8], interchange=[0, 1])
    # CHECK-LABEL: TEST: testTileCompact
    # CHECK: transform.sequence
    # CHECK: %{{.+}}, %{{.+}}:2 = transform.structured.tile_using_for %{{.*}}[4, 8]
    # CHECK: interchange = [0, 1]


@run
@create_sequence
def testTileAttributes(target):
    attr = DenseI64ArrayAttr.get([4, 8])
    ichange = DenseI64ArrayAttr.get([0, 1])
    structured.TileUsingForOp(target, sizes=attr, interchange=ichange)
    # CHECK-LABEL: TEST: testTileAttributes
    # CHECK: transform.sequence
    # CHECK: %{{.+}}, %{{.+}}:2 = transform.structured.tile_using_for %{{.*}}[4, 8]
    # CHECK: interchange = [0, 1]


@run
@create_sequence
def testTileZero(target):
    structured.TileUsingForOp(target, sizes=[4, 0, 2, 0], interchange=[0, 1, 2, 3])
    # CHECK-LABEL: TEST: testTileZero
    # CHECK: transform.sequence
    # CHECK: %{{.+}}, %{{.+}}:2 = transform.structured.tile_using_for %{{.*}}[4, 0, 2, 0]
    # CHECK: interchange = [0, 1, 2, 3]


@run
def testTileDynamic():
    with_pdl = transform_pdl.WithPDLPatternsOp(pdl.OperationType.get())
    with InsertionPoint(with_pdl.body):
        sequence = transform.SequenceOp(
            transform.FailurePropagationMode.Propagate, [], with_pdl.bodyTarget
        )
        with InsertionPoint(sequence.body):
            m1 = transform_pdl.PDLMatchOp(
                pdl.OperationType.get(), sequence.bodyTarget, "first"
            )
            m2 = transform_pdl.PDLMatchOp(
                pdl.OperationType.get(), sequence.bodyTarget, "second"
            )
            structured.TileUsingForOp(sequence.bodyTarget, sizes=[m1, 3, m2, 0])
            transform.YieldOp()
    # CHECK-LABEL: TEST: testTileDynamic
    # CHECK: %[[FIRST:.+]] = pdl_match
    # CHECK: %[[SECOND:.+]] = pdl_match
    # CHECK: %{{.+}}, %{{.+}}:3 = transform.structured.tile_using_for %{{.*}}[%[[FIRST]], 3, %[[SECOND]], 0]


@run
@create_sequence
def testTileExplicitLoopTypeSingle(target):
    structured.TileUsingForOp(
        transform.OperationType.get("scf.for"), target, sizes=[2, 3, 4]
    )
    # CHECK-LABEL: TEST: testTileExplicitLoopTypeSingle
    # CHECK: = transform.structured.tile_using_for %{{.*}} : (!{{.*}}) ->
    # CHECK-COUNT-3: !transform.op<"scf.for">


@run
@create_sequence
def testTileExplicitLoopTypeAll(target):
    types = [
        transform.OperationType.get(x)
        for x in ["scf.for", "scf.parallel", "scf.forall"]
    ]
    structured.TileUsingForOp(types, target, sizes=[2, 3, 4])
    # CHECK-LABEL: TEST: testTileExplicitLoopTypeAll
    # CHECK: = transform.structured.tile
    # CHECK-SAME : (!transform.any_op) -> (!transform.any_op, !transform.op<"scf.for">,
    # CHECK-SAME: !transform.op<"scf.parallel">, !transform.op<"scf.forall">


@run
@create_sequence
def testTileScalable(target):
    structured.TileUsingForOp(
        target,
        sizes=[4, [2]],
    )
    # CHECK-LABEL: TEST: testTileScalable
    # CHECK: transform.sequence
    # CHECK: %{{.+}}, %{{.+}}:2 = transform.structured.tile_using_for %{{.*}}[4, [2]]


@run
@create_sequence
def testTileToForallCompact(target):
    matmul = transform.CastOp(transform.OperationType.get("linalg.matmul"), target)
    structured.TileUsingForallOp(matmul, num_threads=[2, 3, 4])
    # CHECK-LABEL: TEST: testTileToForallCompact
    # CHECK: = transform.structured.tile_using_forall
    # CHECK-SAME: num_threads [2, 3, 4]
    # CHECK-SAME: (!transform.op<"linalg.matmul">) -> (!transform.any_op, !transform.any_op)


@run
@create_sequence
def testTileToForallLoopsAndTileOpTypes(target):
    structured.TileUsingForallOp(
        transform.OperationType.get("scf.forall"),  # loops_type
        transform.OperationType.get("linalg.matmul"),  # tiled_op_type
        target,
        num_threads=[2, 3, 4],
    )
    # CHECK-LABEL: TEST: testTileToForallLoopsAndTileOpTypes
    # CHECK: = transform.structured.tile_using_forall
    # CHECK-SAME: num_threads [2, 3, 4]
    # CHECK-SAME: (!transform.any_op) -> (!transform.op<"scf.forall">, !transform.op<"linalg.matmul">)


@run
@create_sequence
def testTileToForallTileSizes(target):
    structured.TileUsingForallOp(target, tile_sizes=[2, 3, 4])
    # CHECK-LABEL: TEST: testTileToForallTileSizes
    # CHECK: = transform.structured.tile_using_forall
    # CHECK-SAME: tile_sizes [2, 3, 4]


@run
@create_sequence
def testTileToForallMixedDynamic(target):
    n = structured.MatchOp.match_op_names(target, ["test.dummy"])
    structured.TileUsingForallOp(target, num_threads=[n, 3, 4])
    # CHECK-LABEL: TEST: testTileToForallMixedDynamic
    # CHECK: = transform.structured.tile_using_forall
    # CHECK-SAME: num_threads [%{{.*}} : !transform.any_op, 3, 4]


@run
@create_sequence
def testTileToForallPackedDynamic(target):
    n = structured.MatchOp.match_op_names(target, ["test.dummy"])
    structured.TileUsingForallOp(target, num_threads=n)
    # CHECK-LABEL: TEST: testTileToForallPackedDynamic
    # CHECK: = transform.structured.tile_using_forall
    # CHECK-SAME: num_threads *(%0 : !transform.any_op)


@run
@create_sequence
def testTileToForallMapping(target):
    mapping = Attribute.parse("[ #gpu.thread<y>, #gpu.thread<x> ]")
    structured.TileUsingForallOp(target, num_threads=[2, 3], mapping=mapping)
    # CHECK-LABEL: TEST: testTileToForallMapping
    # CHECK: = transform.structured.tile_using_forall
    # CHECK-SAME: mapping = [#gpu.thread<y>, #gpu.thread<x>]


@run
@create_sequence
def testVectorizeChildrenAndApplyPatternsAllAttrs(target):
    structured.VectorizeChildrenAndApplyPatternsOp(
        target,
        disable_multi_reduction_to_contract_patterns=True,
        disable_transfer_permutation_map_lowering_patterns=True,
        vectorize_nd_extract=True,
        vectorize_padding=True,
    )
    # CHECK-LABEL: TEST: testVectorizeChildrenAndApplyPatternsAllAttrs
    # CHECK: transform.sequence
    # CHECK: = transform.structured.vectorize
    # CHECK-SAME: disable_multi_reduction_to_contract_patterns
    # CHECK-SAME: disable_transfer_permutation_map_lowering_patterns
    # CHECK-SAME: vectorize_nd_extract
    # CHECK-SAME: vectorize_padding


@run
@create_sequence
def testVectorizeChildrenAndApplyPatternsNoAttrs(target):
    structured.VectorizeChildrenAndApplyPatternsOp(
        target,
        disable_multi_reduction_to_contract_patterns=False,
        disable_transfer_permutation_map_lowering_patterns=False,
        vectorize_nd_extract=False,
        vectorize_padding=False,
    )
    # CHECK-LABEL: TEST: testVectorizeChildrenAndApplyPatternsNoAttrs
    # CHECK: transform.sequence
    # CHECK: = transform.structured.vectorize
    # CHECK-NOT: disable_multi_reduction_to_contract_patterns
    # CHECK-NOT: disable_transfer_permutation_map_lowering_patterns
    # CHECK-NOT: vectorize_nd_extract
    # CHECK-NOT: vectorize_padding


@run
@create_sequence
def testMatchInterfaceEnum(target):
    names = ArrayAttr.get([StringAttr.get("test.dummy")])
    result_type = transform.AnyOpType.get()
    fused = structured.MatchOp.__base__(
        result_type,
        target,
        ops=names,
        interface=structured.MatchInterfaceEnum.LinalgOp,
    )
    # CHECK-LABEL: TEST: testMatchInterfaceEnum
    # CHECK: transform.sequence
    # CHECK: = transform.structured.match
    # CHECK: interface{LinalgOp}


@run
@create_sequence
def testMatchInterfaceEnumReplaceAttributeBuilder(target):
    @register_attribute_builder("MatchInterfaceEnum", replace=True)
    def match_interface_enum(x, context):
        if x == "LinalgOp":
            y = 0
        elif x == "TilingInterface":
            y = 1
        return IntegerAttr.get(IntegerType.get_signless(32, context=context), y)

    names = ArrayAttr.get([StringAttr.get("test.dummy")])
    result_type = transform.AnyOpType.get()
    fused = structured.MatchOp.__base__(
        result_type,
        target,
        ops=names,
        interface="TilingInterface",
    )
    # CHECK-LABEL: TEST: testMatchInterfaceEnumReplaceAttributeBuilder
    # CHECK: transform.sequence
    # CHECK: = transform.structured.match
    # CHECK: interface{TilingInterface}