# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.
"""The language interface for tl programs."""

from typing import Optional
from tvm.script import tir as T
from tvm.script.parser.tir import *
from tilelang.layout import Layout, Fragment  # noqa: F401
from .parallel import Parallel  # noqa: F401
from .pipeline import Pipelined  # noqa: F401
from .kernel import Kernel  # noqa: F401
from .allocate import (
    alloc_local,  # noqa: F401
    alloc_shared,  # noqa: F401
    alloc_fragment,  # noqa: F401
)
from .copy import copy, c2d_im2col  # noqa: F401
from .gemm import GemmWarpPolicy, gemm  # noqa: F401
from .fill import fill, clear  # noqa: F401
from .reduce import (
    reduce,  # noqa: F401
    reduce_max,  # noqa: F401
    reduce_min,  # noqa: F401
    reduce_sum,  # noqa: F401
    reduce_abssum,  # noqa: F401
)
from .customize import (
    atomic_add,  # noqa: F401
    atomic_addx2,  # noqa: F401
    dp4a,  # noqa: F401
)


def use_swizzle(panel_size: int, order: str = "row", enable: bool = True):
    # If order is row, use rasterization2DRow, otherwise use rasterization2DColumn
    # The panel size is the number of threads in a warp
    # Use to improve the L2 Cache Locality
    device_func = ("rasterization2DRow" if order == "row" else "rasterization2DColumn")
    return T.attr(None, "threadblock_swizzle_pattern",
                  f"tl::{device_func}<{panel_size}>") if enable else None


def annotate_layout(layout_map):
    # layout_map is a dictionary of buffer to layout
    layout_map = {buffer.data: layout for buffer, layout in layout_map.items()}
    return T.block_attr({"layout_map": layout_map})


def import_source(source: Optional[str] = None):
    # source is the source code to be imported
    return T.block_attr({"pragma_import_c": source}) if source is not None else None
