# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
# pylint: disable=invalid-name
"""GEMM kernel generator and profiler for CUTLASS."""

import os
import pickle
from functools import partial

from .gemm_operation import EmitGemmInstance, GemmOperation
from .gemm_profiler import GemmProfilerEmitter
from .gen_tensor_op import EPILOGUE_MAP, GENERATOR_FUNC_TABLE, ProfilerEngine
from .library import (
    DataType,
    DataTypeTag,
    EpilogueFunctor,
    LayoutType,
    SwizzlingFunctor,
    TensorDescription,
)


def create_gemm_operator_with_epilogue(
    op_type,
    tile_description,
    data_type,
    alignment,
    swizzling_functor,
    batched=False,
    layout_b=LayoutType.ColumnMajor,
):
    """
    Instantiate a cutlass kernel from the given configuration,
    along with the epilouge functor
    """
    element_a, element_b, element_c, element_epilogue = data_type

    A = TensorDescription(element_a, LayoutType.RowMajor, alignment)
    B = TensorDescription(element_b, layout_b, alignment)
    C = TensorDescription(element_c, LayoutType.RowMajor, alignment)

    if batched:
        swizzling_functor = SwizzlingFunctor.Batched

    if "residual" in op_type:
        if "hardswish" in op_type:
            activation = "cutlass::epilogue::thread::HardSwish"
        elif "silu" in op_type:
            activation = "cutlass::epilogue::thread::SiLu"
        elif "sigmoid" in op_type:
            activation = "cutlass::epilogue::thread::Sigmoid"
        elif "gelu" in op_type:
            activation = "cutlass::epilogue::thread::GELU"
        elif "relu" in op_type:
            activation = "cutlass::epilogue::thread::ReLu"
        else:
            activation = "cutlass::epilogue::thread::Identity"

        binary_op = "cutlass::multiplies" if "residual_multiply" in op_type else "cutlass::plus"
        unary_op = (
            "cutlass::epilogue::thread::ReLu"
            if op_type.endswith("relu")
            else "cutlass::epilogue::thread::Identity"
        )
        residual_block_info = {
            "activation": activation,
            "binary_op": binary_op,
            "unary_op": unary_op,
        }
        epilogue = EpilogueFunctor.LinearCombinationResidualBlock
        no_beta_scaling = False
    else:
        residual_block_info = None
        epilogue, no_beta_scaling = EPILOGUE_MAP[op_type]

    op = GemmOperation(
        tile_description.minimum_compute_capability,
        tile_description,
        A,
        B,
        C,
        element_epilogue,
        epilogue,
        swizzling_functor,
    )

    return (
        op.procedural_name(),
        EmitGemmInstance().emit(
            op,
            no_beta_scaling=no_beta_scaling,
            batched=batched,
            residual_block_info=residual_block_info,
        ),
    )


def enumerate_gemm_operators(
    tile_descriptions,
    data_type,
    alignment_constraints,
    swizzling_functor=SwizzlingFunctor.Identity8,
    layout_b=LayoutType.ColumnMajor,
):
    """Exhaustively instantiate all kernels from a given configuration."""
    ret = []
    kernel_emitter = EmitGemmInstance()
    profiler_emitter = GemmProfilerEmitter()

    element_a, element_b, element_c, element_epilogue = data_type

    for tile_description in tile_descriptions:
        for alignment in alignment_constraints:
            A = TensorDescription(element_a, LayoutType.RowMajor, alignment)
            B = TensorDescription(element_b, layout_b, alignment)
            C = TensorDescription(element_c, LayoutType.RowMajor, alignment)

            if element_c == DataType.s32 and A.alignment == 1:
                tile_description.threadblock_shape[0] = min(
                    tile_description.threadblock_shape[0], 128
                )
                tile_description.threadblock_shape[1] = min(
                    tile_description.threadblock_shape[1], 128
                )

            op = GemmOperation(
                tile_description.minimum_compute_capability,
                tile_description,
                A,
                B,
                C,
                element_epilogue,
                EpilogueFunctor.LinearCombination,
                swizzling_functor,
            )

            src = profiler_emitter.emit(
                op.procedural_name(),
                kernel_emitter.emit(op, batched=False),
                DataTypeTag[element_a],
                DataTypeTag[element_b],
                DataTypeTag[element_c],
                op.leading_dim(),
            )

            ret.append(
                {
                    "src": src,
                    "op": op,
                    "name": op.procedural_name(),
                    "tile_description": tile_description,
                    "alignment": alignment,
                    "data_type": data_type,
                    "swizzle_functor": swizzling_functor,
                }
            )

    return ret


# TODO(masahi): A sensible way to pick reasonable default kernels
DEFAULT_KERNELS = {
    75: {
        ("float16", "float16"): "cutlass_tensorop_h1688gemm_128x64_32x2_tn_align1",
        ("float16", "float32"): "cutlass_tensorop_s1688gemm_f16_64x64_32x2_tn_align1",
    },
    # align1 variants do not seem to be available for sm80
    80: {
        ("float16", "float16"): "cutlass_tensorop_h1688gemm_128x64_32x2_tn_align1",
        ("float16", "float32"): "cutlass_tensorop_s1688gemm_f16_64x64_32x2_tn_align1",
        # two kernels for tf32 and 3xtf32
        ("float32", "float32"): (
            "cutlass_tensorop_s1688gemm_128x64_32x3_tn_align1",
            "cutlass_tensorop_s1688gemm_64x64_16x3_tn_align1",
        ),
    },
}


class CutlassGemmProfiler:
    """Profile all candidate kernels and select the best one."""

    def __init__(self, sm, cutlass_path, binary_path):
        assert sm in GENERATOR_FUNC_TABLE and sm in DEFAULT_KERNELS, f"sm{sm} not supported yet."
        self.engine = ProfilerEngine(sm, cutlass_path, binary_path)
        self.sm = sm
        self.cache_path = os.path.join(binary_path, "cutlass_gemm_cache.pickle")
        if os.path.exists(self.cache_path):
            self.cache = pickle.load(open(self.cache_path, "rb"))
        else:
            self.cache = {}

    def get_default(
        self,
        op_type,
        out_dtype,
        arg0_dtype,
        arg1_dtype,
        use_3xtf32=True,
        batched=False,
        layout_b=LayoutType.ColumnMajor,
    ):
        """Return the default kernel for the requested architecture.
        For now, the default kernel was picked arbitrary.
        """
        ops = GENERATOR_FUNC_TABLE[self.sm](
            out_dtype,
            arg0_dtype,
            arg1_dtype,
            partial(enumerate_gemm_operators, layout_b=layout_b),
            lambda align: align == 1,  # Only request align1 kernels
            use_3xtf32,
            profile_all_alignments=True,  # To include all align1 kernels
            # TODO(masahi): Invesitigate when fp32 accumulation is needed for gemm
            accumlator_dtype=out_dtype,
        )

        default_kernel_name = DEFAULT_KERNELS[self.sm][(arg0_dtype, out_dtype)]

        if arg0_dtype == "float32":
            default_kernel_name = (
                default_kernel_name[0] if not use_3xtf32 else default_kernel_name[1]
            )

        filtered = list(filter(lambda op: op["name"] == default_kernel_name, ops))
        assert len(filtered) == 1
        op = filtered[0]
        name, opdef = create_gemm_operator_with_epilogue(
            op_type,
            op["tile_description"],
            op["data_type"],
            op["alignment"],
            op["swizzle_functor"],
            batched=batched,
            layout_b=layout_b,
        )
        op.update({"name": name, "opdef": opdef})
        return op

    def select_op(
        self,
        M,
        N,
        K,
        out_dtype,
        arg0_dtype,
        arg1_dtype,
        use_3xtf32,
        profile_all_alignments=False,
        find_first_valid=False,
        use_multiprocessing=False,
        layout_b=LayoutType.ColumnMajor,
    ):
        """
        Profile and select the best kernel from candidate kernels.
        See the documentation for the profile method below.
        """
        if (M, N, K) in self.cache:
            op = self.cache[(M, N, K)]
            return op

        # TODO(masahi): CUTLASS alignment check on gemm kernels is too restrictive.
        # See https://github.com/NVIDIA/cutlass/issues/362.
        # When the above issue is resolved, we can remove the alignment check on M below.

        ops = GENERATOR_FUNC_TABLE[self.sm](
            out_dtype,
            arg0_dtype,
            arg1_dtype,
            partial(enumerate_gemm_operators, layout_b=layout_b),
            lambda align: all([dim % align == 0 for dim in [M, N, K]]),
            use_3xtf32,
            profile_all_alignments=profile_all_alignments,
            # TODO(masahi): Invesitigate when fp32 accumulation is needed for gemm
            accumlator_dtype=out_dtype,
        )

        if not find_first_valid:
            self.engine.compile_all(ops, use_multiprocessing)

        for op in ops:
            out = self.engine.evaluate(op, [M, N, K])
            op["runtime"] = out
            if out < float("inf") and find_first_valid:
                self.cache[(M, N, K)] = op
                return op

        op = min(ops, key=lambda i: i["runtime"])
        self.cache[(M, N, K)] = op
        with open(self.cache_path, "wb") as f:
            pickle.dump(self.cache, f)
        return op

    def profile(
        self,
        op_type,
        M,
        N,
        K,
        out_dtype,
        arg0_dtype,
        arg1_dtype,
        use_3xtf32=True,
        profile_all_alignments=False,
        find_first_valid=False,
        use_multiprocessing=False,
        batched=False,
        layout_b=LayoutType.ColumnMajor,
    ):
        """Profile and select the best kernel from candidate kernels.
        If find_first_valid is True, return immediately after the first applicable kernel is found.
        If use_multiprocessing is True, compile all profiler executables in parallel.
        """
        op = self.select_op(
            M,
            N,
            K,
            out_dtype,
            arg0_dtype,
            arg1_dtype,
            use_3xtf32,
            profile_all_alignments=profile_all_alignments,
            find_first_valid=find_first_valid,
            use_multiprocessing=use_multiprocessing,
            layout_b=layout_b,
        )

        name, opdef = create_gemm_operator_with_epilogue(
            op_type,
            op["tile_description"],
            op["data_type"],
            op["alignment"],
            op["swizzle_functor"],
            batched=batched,
            layout_b=layout_b,
        )

        return name, opdef, op["runtime"]
