# Copyright 2021 Yan Yan
#
# Licensed 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.
import os
from enum import Enum
from cumm.gemm.main import gen_shuffle_params_v2 as gen_shuffle_params, GemmAlgoParams
from cumm.gemm import kernel
from typing import List
from cumm.gemm.algospec.core import TensorOp
from cumm.conv.main import gen_gemm_params as gen_conv_params, ConvFwdAndBwdInput, ConvBwdWeight, ConvIterAlgo, GemmAlgo
from cumm.conv.bases import (NCHW, NHWC, ConvIterAlgo, ConvLayout,
                             ConvLayoutType, ConvMode, ConvOpType)
from spconv.constants import NDIM_DONT_CARE, SPCONV_INT8_DEBUG


class ConvAlgo(Enum):
    Native = 0
    MaskImplicitGemm = 1
    MaskSplitImplicitGemm = 2


class AlgoHint(Enum):
    NoHint = 0b000
    Fowrard = 0b001
    BackwardInput = 0b010
    BackwardWeight = 0b100


# we can't add more kernels here because build in github action is very slow.
# TODO two step build: build gemm kernels first, then bind for every python

SHUFFLE_SIMT_PARAMS: List[GemmAlgoParams] = [
    # *gen_shuffle_params((64, 128, 32), (32, 64, 32), ["s8,s8,s8,s32,s32"], "",
    #                     2, kernel.GemmAlgo.SimtDP4A, None),
    # *gen_shuffle_params((128, 64, 32), (64, 32, 32), ["s8,s8,s8,s32,s32"], "",
    #                     2, kernel.GemmAlgo.SimtDP4A, None),
    # *gen_shuffle_params((128, 128, 32), (32, 64, 32), ["s8,s8,s8,s32,s32"],
    #                     "", 2, kernel.GemmAlgo.SimtDP4A, None),
    # *gen_shuffle_params(
    #     (128, 128, 32),
    #     (64, 32, 32), ["s8,s8,s8,s32,s32"], "", 2,
    #     kernel.GemmAlgo.SimtDP4A, None),
    # *gen_shuffle_params((64, 64, 32), (32, 32, 32), ["s8,s8,s8,s32,s32"], "",
    #                     2, kernel.GemmAlgo.SimtDP4A, None),
    *gen_shuffle_params((64, 256, 8), (32, 64, 8), ["f32,f32,f32,f32,f32"],
                        "f32,f32,f32,f32,f32", 2, kernel.GemmAlgo.Simt, None),
    # *gen_shuffle_params(
    #     (64, 256, 8),
    #     (64, 32, 8), ["f32,f32,f32,f32,f32"], 2, kernel.GemmAlgo.Simt, None),
    *gen_shuffle_params((32, 128, 16), (32, 32, 8), ["f32,f32,f32,f32,f32"],
                        "f32,f32,f32,f32,f32", 2, kernel.GemmAlgo.Simt, None),
    *gen_shuffle_params((32, 512, 8), (32, 64, 8), ["f32,f32,f32,f32,f32"],
                        "f32,f32,f32,f32,f32", 2, kernel.GemmAlgo.Simt, None),
    # *gen_shuffle_params(
    #     (128, 128, 8),
    #     (64, 32, 8), ["f32,f32,f32,f32,f32"], 2, kernel.GemmAlgo.Simt, None),
    *gen_shuffle_params((128, 128, 8), (32, 64, 8), ["f32,f32,f32,f32,f32"],
                        "f32,f32,f32,f32,f32", 2, kernel.GemmAlgo.Simt, None),
    *gen_shuffle_params((64, 128, 8), (32, 64, 8), ["f32,f32,f32,f32,f32"],
                        "f32,f32,f32,f32,f32", 2, kernel.GemmAlgo.Simt, None),
    # *gen_shuffle_params(
    #     (64, 128, 8),
    #     (64, 32, 8), ["f32,f32,f32,f32,f32"], 2, kernel.GemmAlgo.Simt, None),
    # *gen_shuffle_params(
    #     (128, 64, 8),
    #     (32, 64, 8), ["f32,f32,f32,f32,f32"], 2, kernel.GemmAlgo.Simt, None),
    *gen_shuffle_params((128, 64, 8), (64, 32, 8), ["f32,f32,f32,f32,f32"],
                        "f32,f32,f32,f32,f32", 2, kernel.GemmAlgo.Simt, None),
    *gen_shuffle_params((64, 64, 8), (32, 32, 8), ["f32,f32,f32,f32,f32"],
                        "f32,f32,f32,f32,f32", 2, kernel.GemmAlgo.Simt, None),
    *gen_shuffle_params((32, 64, 16), (32, 32, 8), ["f32,f32,f32,f32,f32"],
                        "f32,f32,f32,f32,f32", 2, kernel.GemmAlgo.Simt, None),
    *gen_shuffle_params((64, 32, 16), (32, 32, 8), ["f32,f32,f32,f32,f32"],
                        "f32,f32,f32,f32,f32", 2, kernel.GemmAlgo.Simt, None),
    *gen_shuffle_params((32, 32, 32), (32, 32, 8), ["f32,f32,f32,f32,f32"],
                        "f32,f32,f32,f32,f32", 2, kernel.GemmAlgo.Simt, None),
    # fall back kernels if mat is misaligned for half
    # TODO use access-per-vector kernel instead of simt kernel for fallback
    *gen_shuffle_params((128, 128, 8), (32, 64, 8), ["f16,f16,f16,f32,f32"],
                        "f16,f16,f16,f32,f32", 2, kernel.GemmAlgo.Simt, None),
    *gen_shuffle_params((32, 64, 32), (32, 32, 8), ["f16,f16,f16,f32,f32"],
                        "f16,f16,f16,f32,f32", 2, kernel.GemmAlgo.Simt, None),
    *gen_shuffle_params((32, 32, 32), (32, 32, 8), ["f16,f16,f16,f32,f32"],
                        "f16,f16,f16,f32,f32", 2, kernel.GemmAlgo.Simt, None),
    # *gen_shuffle_params(
    #     (64, 64, 16),
    #     (32, 32, 8), ["f16,f16,f16,f16,f16"], "f16,f16,f16,f32,f32", 2, kernel.GemmAlgo.Simt, None),
    *gen_shuffle_params((64, 128, 16), (32, 64, 8), ["f16,f16,f16,f32,f32"],
                        "f16,f16,f16,f32,f32", 2, kernel.GemmAlgo.Simt, None),
    *gen_shuffle_params((64, 64, 8), (32, 32, 8), ["f16,f16,f16,f32,f32"],
                        "f16,f16,f16,f32,f32", 2, kernel.GemmAlgo.Simt, None),
]

SHUFFLE_VOLTA_PARAMS: List[GemmAlgoParams] = [
    *gen_shuffle_params(
        (64, 64, 32),
        (32, 32, 32), ["f16,f16,f16,f16,f16"], "f16,f16,f16,f32,f32", 2,
        kernel.GemmAlgo.Volta, TensorOp((8, 8, 4))),
    # *gen_shuffle_params(
    #     (128, 128, 32),
    #     (64, 64, 32), ["f16,f16,f16,f16,f16"], "f16,f16,f16,f32,f32", 2,
    #     kernel.GemmAlgo.Volta, TensorOp((8, 8, 4))),
    *gen_shuffle_params(
        (128, 256, 32),
        (64, 64, 32), ["f16,f16,f16,f16,f16"], "f16,f16,f16,f32,f32", 2,
        kernel.GemmAlgo.Volta, TensorOp((8, 8, 4))),
    *gen_shuffle_params(
        (256, 128, 32),
        (64, 64, 32), ["f16,f16,f16,f16,f16"], "f16,f16,f16,f32,f32", 2,
        kernel.GemmAlgo.Volta, TensorOp((8, 8, 4))),
    *gen_shuffle_params(
        (128, 64, 32),
        (64, 32, 32), ["f16,f16,f16,f16,f16"], "f16,f16,f16,f32,f32", 2,
        kernel.GemmAlgo.Volta, TensorOp((8, 8, 4))),
    *gen_shuffle_params(
        (64, 128, 32),
        (32, 64, 32), ["f16,f16,f16,f16,f16"], "f16,f16,f16,f32,f32", 2,
        kernel.GemmAlgo.Volta, TensorOp((8, 8, 4))),
]
# SHUFFLE_VOLTA_PARAMS = []
SHUFFLE_TURING_PARAMS: List[GemmAlgoParams] = [
    *gen_shuffle_params(
        (64, 64, 32),
        (32, 32, 32), ["f16,f16,f16,f16,f16"], "f16,f16,f16,f32,f32", 2,
        kernel.GemmAlgo.Turing, TensorOp((16, 8, 8))),
    *gen_shuffle_params(
        (128, 128, 32),
        (32, 64, 32), ["f16,f16,f16,f16,f16"], "f16,f16,f16,f32,f32", 2,
        kernel.GemmAlgo.Turing, TensorOp((16, 8, 8))),
    # *gen_shuffle_params(
    #     (128, 128, 32),
    #     (64, 32, 32), ["f16,f16,f16,f16,f16"], "f16,f16,f16,f32,f32", 2,
    #     kernel.GemmAlgo.Turing, TensorOp((16, 8, 8))),
    *gen_shuffle_params(
        (64, 64, 64),
        (32, 32, 32), ["f16,f16,f16,f16,f16"], "f16,f16,f16,f32,f32", 2,
        kernel.GemmAlgo.Turing, TensorOp((16, 8, 8))),
    *gen_shuffle_params(
        (64, 128, 64),
        (32, 64, 32), ["f16,f16,f16,f16,f16"], "f16,f16,f16,f32,f32", 2,
        kernel.GemmAlgo.Turing, TensorOp((16, 8, 8))),
    *gen_shuffle_params(
        (128, 256, 32),
        (64, 64, 32), ["f16,f16,f16,f16,f16"], "f16,f16,f16,f32,f32", 2,
        kernel.GemmAlgo.Turing, TensorOp((16, 8, 8))),
    *gen_shuffle_params(
        (256, 128, 32),
        (64, 64, 32), ["f16,f16,f16,f16,f16"], "f16,f16,f16,f32,f32", 2,
        kernel.GemmAlgo.Turing, TensorOp((16, 8, 8))),
    *gen_shuffle_params(
        (128, 64, 32),
        (64, 32, 32), ["f16,f16,f16,f16,f16"], "f16,f16,f16,f32,f32", 2,
        kernel.GemmAlgo.Turing, TensorOp((16, 8, 8))),
    *gen_shuffle_params(
        (64, 128, 32),
        (32, 64, 32), ["f16,f16,f16,f16,f16"], "f16,f16,f16,f32,f32", 2,
        kernel.GemmAlgo.Turing, TensorOp((16, 8, 8))),
    # *gen_shuffle_params((64, 64, 32), (32, 32, 32), ["s8,s8,s8,s32,s32"], "",
    #                     2, kernel.GemmAlgo.Turing, TensorOp((8, 8, 16))),
    # *gen_shuffle_params(
    #     (128, 128, 32),
    #     (32, 64, 32), ["s8,s8,s8,s32,s32"], "", 2, kernel.GemmAlgo.Turing,
    #     TensorOp((8, 8, 16))),
    # # *gen_shuffle_params(
    # #     (128, 128, 32),
    # #     (64, 32, 32), ["s8,s8,s8,s32,s32", "s8,s8,s32,s32,s32"], "", 2,
    # #     kernel.GemmAlgo.Turing, TensorOp((8, 8, 16))),
    # *gen_shuffle_params(
    #     (128, 256, 32),
    #     (64, 64, 32), ["s8,s8,s8,s32,s32"], "", 2, kernel.GemmAlgo.Turing,
    #     TensorOp((8, 8, 16))),
    # *gen_shuffle_params(
    #     (256, 128, 32),
    #     (64, 64, 32), ["s8,s8,s8,s32,s32"], "", 2, kernel.GemmAlgo.Turing,
    #     TensorOp((8, 8, 16))),
    # *gen_shuffle_params((128, 64, 32), (64, 32, 32), ["s8,s8,s8,s32,s32"], "",
    #                     2, kernel.GemmAlgo.Turing, TensorOp((8, 8, 16))),
    # *gen_shuffle_params((64, 128, 32), (32, 64, 32), ["s8,s8,s8,s32,s32"], "",
    #                     2, kernel.GemmAlgo.Turing, TensorOp((8, 8, 16))),
]

SHUFFLE_AMPERE_PARAMS = [
    # *gen_shuffle_params(
    #     (128, 128, 64),
    #     (64, 64, 64), ["s8,s8,s8,s32,s32"], "", 3, kernel.GemmAlgo.Ampere,
    #     TensorOp((8, 8, 16))),
    # *gen_shuffle_params(
    #     (128, 64, 64),
    #     (64, 32, 64), ["s8,s8,s8,s32,s32"], "", 3, kernel.GemmAlgo.Ampere,
    #     TensorOp((8, 8, 16))),
]

# SHUFFLE_TURING_PARAMS = []
# here we must use f32 for simt f16 accumulators because
# half intristics is VERY SLOW in GTX 1000 series.
IMPLGEMM_SIMT_PARAMS = [
    *gen_conv_params(ConvFwdAndBwdInput, (32, 128, 16), (32, 32, 8),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     2, ["f32,f32,f32,f32,f32", "f16,f16,f16,f32,f32"],
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Simt,
                     None,
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=1),
    *gen_conv_params(ConvFwdAndBwdInput, (32, 256, 8), (32, 64, 8),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     2, ["f32,f32,f32,f32,f32", "f16,f16,f16,f32,f32"],
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Simt,
                     None,
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=1),
    *gen_conv_params(ConvFwdAndBwdInput, (32, 64, 16), (32, 32, 8),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     2, ["f32,f32,f32,f32,f32", "f16,f16,f16,f32,f32"],
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Simt,
                     None,
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=1),
    *gen_conv_params(ConvFwdAndBwdInput, (32, 32, 32), (32, 32, 8),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     2, ["f32,f32,f32,f32,f32", "f16,f16,f16,f32,f32"],
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Simt,
                     None,
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=1),
    *gen_conv_params(ConvFwdAndBwdInput, (64, 256, 8), (32, 64, 8),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     2, ["f32,f32,f32,f32,f32", "f16,f16,f16,f32,f32"],
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Simt,
                     None,
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=1),
    *gen_conv_params(ConvFwdAndBwdInput, (64, 128, 8), (32, 64, 8),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     2, ["f32,f32,f32,f32,f32", "f16,f16,f16,f32,f32"],
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Simt,
                     None,
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=1),
    *gen_conv_params(ConvFwdAndBwdInput, (64, 64, 8), (32, 32, 8),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     2, ["f32,f32,f32,f32,f32", "f16,f16,f16,f32,f32"],
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Simt,
                     None,
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=1),
    *gen_conv_params(ConvFwdAndBwdInput, (64, 32, 16), (32, 32, 8),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     2, ["f32,f32,f32,f32,f32", "f16,f16,f16,f32,f32"],
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Simt,
                     None,
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=1),
    *gen_conv_params(ConvBwdWeight, (32, 128, 16), (32, 32, 8),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     2, ["f32,f32,f32,f32,f32", "f16,f16,f16,f32,f32"],
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Simt,
                     None,
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=1),
    # *gen_conv_params(ConvBwdWeight, (32, 256, 8), (32, 64, 8), NDIM_DONT_CARE, ConvIterAlgo.Optimized, 2, ["f32,f32,f32,f32,f32"],
    #     NHWC, NHWC, NHWC, GemmAlgo.Simt, None, mask_sparse=True, increment_k_first=True, access_per_vector=1),
    *gen_conv_params(ConvBwdWeight, (32, 64, 16), (32, 32, 8),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     2, ["f32,f32,f32,f32,f32", "f16,f16,f16,f32,f32"],
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Simt,
                     None,
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=1),
    *gen_conv_params(ConvBwdWeight, (32, 32, 32), (32, 32, 8),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     2, ["f32,f32,f32,f32,f32", "f16,f16,f16,f32,f32"],
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Simt,
                     None,
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=1),
    *gen_conv_params(ConvBwdWeight, (64, 256, 8), (32, 64, 8),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     2, ["f32,f32,f32,f32,f32", "f16,f16,f16,f32,f32"],
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Simt,
                     None,
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=1),
    *gen_conv_params(ConvBwdWeight, (64, 128, 8), (32, 64, 8),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     2, ["f32,f32,f32,f32,f32", "f16,f16,f16,f32,f32"],
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Simt,
                     None,
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=1),
    *gen_conv_params(ConvBwdWeight, (64, 64, 8), (32, 32, 8),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     2, ["f32,f32,f32,f32,f32", "f16,f16,f16,f32,f32"],
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Simt,
                     None,
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=1),
    *gen_conv_params(ConvBwdWeight, (64, 32, 16), (32, 32, 8),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     2, ["f32,f32,f32,f32,f32", "f16,f16,f16,f32,f32"],
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Simt,
                     None,
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=1),
    *gen_conv_params(ConvBwdWeight, (128, 128, 8), (32, 64, 8),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     2, ["f32,f32,f32,f32,f32", "f16,f16,f16,f32,f32"],
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Simt,
                     None,
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=1),
    *gen_conv_params(ConvBwdWeight, (128, 64, 8), (64, 32, 8),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     2, ["f32,f32,f32,f32,f32", "f16,f16,f16,f32,f32"],
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Simt,
                     None,
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=1),
]

IMPLGEMM_VOLTA_PARAMS = [
    *gen_conv_params(ConvFwdAndBwdInput, (64, 64, 32), (32, 32, 32),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     2, ["f16,f16,f16,f16,f16", "f16,f16,f16,f32,f32"],
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Volta,
                     TensorOp((8, 8, 4)),
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=1),
    *gen_conv_params(ConvFwdAndBwdInput, (64, 64, 32), (32, 32, 32),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     2, ["f16,f16,f16,f16,f16", "f16,f16,f16,f32,f32"],
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Volta,
                     TensorOp((8, 8, 4)),
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=0),
    *gen_conv_params(ConvFwdAndBwdInput, (64, 128, 32), (32, 64, 32),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     2, ["f16,f16,f16,f16,f16", "f16,f16,f16,f32,f32"],
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Volta,
                     TensorOp((8, 8, 4)),
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=1),
    *gen_conv_params(ConvFwdAndBwdInput, (32, 256, 32), (32, 64, 32),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     2, ["f16,f16,f16,f16,f16", "f16,f16,f16,f32,f32"],
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Volta,
                     TensorOp((8, 8, 4)),
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=1),
    *gen_conv_params(ConvBwdWeight, (64, 64, 32), (32, 32, 32),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     2, ["f16,f16,f16,f32,f32"],
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Volta,
                     TensorOp((8, 8, 4)),
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=1),
    *gen_conv_params(ConvBwdWeight, (64, 64, 32), (32, 32, 32),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     2, ["f16,f16,f16,f32,f32"],
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Volta,
                     TensorOp((8, 8, 4)),
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=0),
    *gen_conv_params(ConvBwdWeight, (128, 128, 32), (32, 64, 32),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     2, ["f16,f16,f16,f32,f32"],
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Volta,
                     TensorOp((8, 8, 4)),
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=1),
]


IMPLGEMM_AMPERE_PARAMS = [

    *gen_conv_params(ConvFwdAndBwdInput, (32, 32, 16), (16, 16, 16),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     2, ["f32,f32,f32,f32,f32"],
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Ampere,
                     TensorOp((16, 8, 8), top_dtypes="tf32,tf32,f32"),
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=0),
    *gen_conv_params(ConvFwdAndBwdInput, (32, 32, 16), (16, 16, 16),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     [2, 3, 4], ["f32,f32,f32,f32,f32"],
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Ampere,
                     TensorOp((16, 8, 8), top_dtypes="tf32,tf32,f32"),
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=1),

    *gen_conv_params(ConvFwdAndBwdInput, (64, 64, 32), (32, 32, 32),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     [2, 3, 4], ["f32,f32,f32,f32,f32"],
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Ampere,
                     TensorOp((16, 8, 8), top_dtypes="tf32,tf32,f32"),
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=1),

    *gen_conv_params(ConvFwdAndBwdInput, (64, 128, 32), (32, 64, 32),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     [2, 3, 4], ["f16,f16,f16,f16,f16", "f16,f16,f16,f32,f32"],
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Ampere,
                     TensorOp((16, 8, 8)),
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=1),

    *gen_conv_params(ConvFwdAndBwdInput, (64, 64, 32), (32, 32, 32),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     [2, 3, 4], ["f16,f16,f16,f16,f16", "f16,f16,f16,f32,f32"],
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Ampere,
                     TensorOp((16, 8, 8)),
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=1),

    *gen_conv_params(ConvFwdAndBwdInput, (128, 64, 32), (64, 32, 32),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     [2, 3, 4], ["f16,f16,f16,f16,f16", "f16,f16,f16,f32,f32"],
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Ampere,
                     TensorOp((16, 8, 8)),
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=1),
    
    *gen_conv_params(ConvBwdWeight, (64, 64, 32), (32, 32, 32),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     [2, 3, 4],
                     "f32,f32,f32,f32,f32",
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Ampere,
                     TensorOp((16, 8, 8), top_dtypes="tf32,tf32,f32"),
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=1),

    *gen_conv_params(ConvBwdWeight, (64, 128, 32), (32, 64, 32),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     [2, 3, 4],
                     "f16,f16,f16,f32,f32",
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Ampere,
                     TensorOp((16, 8, 8)),
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=1),

    *gen_conv_params(ConvBwdWeight, (64, 64, 32), (32, 32, 32),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     [2, 3, 4, 5],
                     "f16,f16,f16,f32,f32",
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Ampere,
                     TensorOp((16, 8, 8)),
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=1),
    *gen_conv_params(ConvFwdAndBwdInput, (64, 64, 32), (32, 32, 32),
                    NDIM_DONT_CARE,
                    ConvIterAlgo.Optimized,
                    [2, 3, 4],
                    ["s8,s8,s8,s32,f32", "s8,s8,s8,s32,f16", "s8,s8,f32,s32,f32", "s8,s8,f32,s32,f16", "s8,s8,f16,s32,f32", "s8,s8,f16,s32,f16"],
                    NHWC,
                    NHWC,
                    NHWC,
                    GemmAlgo.Ampere,
                    TensorOp((16, 8, 16)),
                    mask_sparse=True,
                    increment_k_first=True,
                    access_per_vector=1,
                    is_nvrtc=True,
                    int8_inference=True),
]

if not SPCONV_INT8_DEBUG:
    IMPLGEMM_AMPERE_PARAMS.extend([ 
        *gen_conv_params(ConvFwdAndBwdInput, (32, 32, 32), (32, 32, 32),
                        NDIM_DONT_CARE,
                        ConvIterAlgo.Optimized,
                        [2, 3, 4],
                        ["s8,s8,s8,s32,f32", "s8,s8,s8,s32,f16"],
                        NHWC,
                        NHWC,
                        NHWC,
                        GemmAlgo.Ampere,
                        TensorOp((16, 8, 16)),
                        mask_sparse=True,
                        increment_k_first=True,
                        access_per_vector=1,
                        is_nvrtc=True,
                        int8_inference=True),
        *gen_conv_params(ConvFwdAndBwdInput, (64, 32, 32), (32, 32, 32),
                        NDIM_DONT_CARE,
                        ConvIterAlgo.Optimized,
                        [2, 3, 4],
                        ["s8,s8,s8,s32,f32", "s8,s8,s8,s32,f16"],
                        NHWC,
                        NHWC,
                        NHWC,
                        GemmAlgo.Ampere,
                        TensorOp((16, 8, 16)),
                        mask_sparse=True,
                        increment_k_first=True,
                        access_per_vector=1,
                        is_nvrtc=True,
                        int8_inference=True),
        *gen_conv_params(ConvFwdAndBwdInput, (32, 64, 32), (32, 32, 32),
                        NDIM_DONT_CARE,
                        ConvIterAlgo.Optimized,
                        [2, 3, 4],
                        ["s8,s8,s8,s32,f32", "s8,s8,s8,s32,f16"],
                        NHWC,
                        NHWC,
                        NHWC,
                        GemmAlgo.Ampere,
                        TensorOp((16, 8, 16)),
                        mask_sparse=True,
                        increment_k_first=True,
                        access_per_vector=1,
                        is_nvrtc=True,
                        int8_inference=True),

        *gen_conv_params(ConvFwdAndBwdInput, (128, 64, 64), (64, 32, 64),
                        NDIM_DONT_CARE,
                        ConvIterAlgo.Optimized,
                        [2, 3, 4],
                        ["s8,s8,s8,s32,f32", "s8,s8,s8,s32,f16"],
                        NHWC,
                        NHWC,
                        NHWC,
                        GemmAlgo.Ampere,
                        TensorOp((16, 8, 32)),
                        mask_sparse=True,
                        increment_k_first=True,
                        access_per_vector=1,
                        is_nvrtc=True,
                        int8_inference=True),

        *gen_conv_params(ConvFwdAndBwdInput, (128, 64, 32), (64, 32, 32),
                        NDIM_DONT_CARE,
                        ConvIterAlgo.Optimized,
                        [2, 3, 4],
                        ["s8,s8,s8,s32,f32", "s8,s8,s8,s32,f16"],
                        NHWC,
                        NHWC,
                        NHWC,
                        GemmAlgo.Ampere,
                        TensorOp((16, 8, 16)),
                        mask_sparse=True,
                        increment_k_first=True,
                        access_per_vector=1,
                        is_nvrtc=True,
                        int8_inference=True),
        
        *gen_conv_params(ConvFwdAndBwdInput, (64, 128, 64), (32, 64, 64),
                        NDIM_DONT_CARE,
                        ConvIterAlgo.Optimized,
                        [2, 3, 4],
                        ["s8,s8,s8,s32,f32", "s8,s8,s8,s32,f16"],
                        NHWC,
                        NHWC,
                        NHWC,
                        GemmAlgo.Ampere,
                        TensorOp((16, 8, 32)),
                        mask_sparse=True,
                        increment_k_first=True,
                        access_per_vector=1,
                        is_nvrtc=True,
                        int8_inference=True),

        *gen_conv_params(ConvFwdAndBwdInput, (64, 128, 32), (32, 64, 32),
                        NDIM_DONT_CARE,
                        ConvIterAlgo.Optimized,
                        [2, 3, 4],
                        ["s8,s8,s8,s32,f32", "s8,s8,s8,s32,f16"],
                        NHWC,
                        NHWC,
                        NHWC,
                        GemmAlgo.Ampere,
                        TensorOp((16, 8, 16)),
                        mask_sparse=True,
                        increment_k_first=True,
                        access_per_vector=1,
                        is_nvrtc=True,
                        int8_inference=True),
            
        *gen_conv_params(ConvFwdAndBwdInput, (64, 64, 64), (32, 32, 64),
                        NDIM_DONT_CARE,
                        ConvIterAlgo.Optimized,
                        [2, 3, 4],
                        ["s8,s8,s8,s32,f32", "s8,s8,s8,s32,f16"],
                        NHWC,
                        NHWC,
                        NHWC,
                        GemmAlgo.Ampere,
                        TensorOp((16, 8, 32)),
                        mask_sparse=True,
                        increment_k_first=True,
                        access_per_vector=1,
                        is_nvrtc=True,
                        int8_inference=True),
        
        *gen_conv_params(ConvFwdAndBwdInput, (128, 128, 64), (64, 64, 64),
                        NDIM_DONT_CARE,
                        ConvIterAlgo.Optimized,
                        [2, 3, 4],
                        ["s8,s8,s8,s32,f32", "s8,s8,s8,s32,f16"],
                        NHWC,
                        NHWC,
                        NHWC,
                        GemmAlgo.Ampere,
                        TensorOp((16, 8, 32)),
                        mask_sparse=True,
                        increment_k_first=True,
                        access_per_vector=1,
                        is_nvrtc=True,
                        int8_inference=True),
        
        *gen_conv_params(ConvFwdAndBwdInput, (128, 256, 64), (64, 128, 64),
                        NDIM_DONT_CARE,
                        ConvIterAlgo.Optimized,
                        [2, 3, 4],
                        ["s8,s8,s8,s32,f32", "s8,s8,s8,s32,f16"],
                        NHWC,
                        NHWC,
                        NHWC,
                        GemmAlgo.Ampere,
                        TensorOp((16, 8, 32)),
                        mask_sparse=True,
                        increment_k_first=True,
                        access_per_vector=1,
                        is_nvrtc=True,
                        int8_inference=True),
        
        *gen_conv_params(ConvFwdAndBwdInput, (256, 128, 64), (128, 64, 64),
                        NDIM_DONT_CARE,
                        ConvIterAlgo.Optimized,
                        [2, 3, 4],
                        ["s8,s8,s8,s32,f32", "s8,s8,s8,s32,f16"],
                        NHWC,
                        NHWC,
                        NHWC,
                        GemmAlgo.Ampere,
                        TensorOp((16, 8, 32)),
                        mask_sparse=True,
                        increment_k_first=True,
                        access_per_vector=1,
                        is_nvrtc=True,
                        int8_inference=True),
        
        *gen_conv_params(ConvFwdAndBwdInput, (128, 128, 128), (64, 64, 128),
                        NDIM_DONT_CARE,
                        ConvIterAlgo.Optimized,
                        [2, 3],
                        ["s8,s8,s8,s32,f32", "s8,s8,s8,s32,f16"],
                        NHWC,
                        NHWC,
                        NHWC,
                        GemmAlgo.Ampere,
                        TensorOp((16, 8, 32)),
                        mask_sparse=True,
                        increment_k_first=True,
                        access_per_vector=1,
                        is_nvrtc=True,
                        int8_inference=True),
    ])

IMPLGEMM_TURING_PARAMS = [

    *gen_conv_params(ConvFwdAndBwdInput, (64, 64, 32), (32, 32, 32),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     2,
                     ["s8,s8,s8,s32,f32", "s8,s8,s8,s32,f16", "s8,s8,f32,s32,f32", "s8,s8,f32,s32,f16", "s8,s8,f16,s32,f32", "s8,s8,f16,s32,f16"],
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Turing,
                     TensorOp((8, 8, 16)),
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=1,
                     is_nvrtc=True,
                     int8_inference=True),
    *gen_conv_params(ConvFwdAndBwdInput, (64, 64, 32), (32, 32, 32),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     2,
                     ["s8,s8,s8,s32,f32", "s8,s8,s8,s32,f16", "s8,s8,f32,s32,f32", "s8,s8,f32,s32,f16", "s8,s8,f16,s32,f32", "s8,s8,f16,s32,f16"],
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Turing,
                     TensorOp((8, 8, 16)),
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=0,
                     is_nvrtc=True,
                     int8_inference=True),
    *gen_conv_params(ConvFwdAndBwdInput, (32, 16, 16), (16, 16, 16),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     2, ["f16,f16,f16,f16,f16", "f16,f16,f16,f32,f32"],
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Turing,
                     TensorOp((16, 8, 8)),
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=0),
    *gen_conv_params(ConvFwdAndBwdInput, (32, 16, 16), (16, 16, 16),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     2, ["f16,f16,f16,f16,f16", "f16,f16,f16,f32,f32"],
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Turing,
                     TensorOp((16, 8, 8)),
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=1),    
    *gen_conv_params(ConvFwdAndBwdInput, (64, 32, 16), (32, 16, 16),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     2, ["f16,f16,f16,f16,f16", "f16,f16,f16,f32,f32"],
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Turing,
                     TensorOp((16, 8, 8)),
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=0),
    *gen_conv_params(ConvFwdAndBwdInput, (64, 32, 16), (32, 16, 16),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     2, ["f16,f16,f16,f16,f16", "f16,f16,f16,f32,f32"],
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Turing,
                     TensorOp((16, 8, 8)),
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=1),

    *gen_conv_params(ConvFwdAndBwdInput, (32, 64, 32), (32, 32, 16),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     2, ["f16,f16,f16,f16,f16", "f16,f16,f16,f32,f32"],
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Turing,
                     TensorOp((16, 8, 8)),
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=0),
    *gen_conv_params(ConvFwdAndBwdInput, (32, 64, 32), (32, 32, 16),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     2, ["f16,f16,f16,f16,f16", "f16,f16,f16,f32,f32"],
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Turing,
                     TensorOp((16, 8, 8)),
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=1),
    *gen_conv_params(ConvFwdAndBwdInput, (32, 256, 32), (32, 64, 32),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     2, ["f16,f16,f16,f16,f16", "f16,f16,f16,f32,f32"],
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Turing,
                     TensorOp((16, 8, 8)),
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=1),
    *gen_conv_params(ConvFwdAndBwdInput, (32, 128, 32), (32, 32, 32),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     2, ["f16,f16,f16,f16,f16", "f16,f16,f16,f32,f32"],
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Turing,
                     TensorOp((16, 8, 8)),
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=1),
    *gen_conv_params(ConvFwdAndBwdInput, (32, 128, 64), (32, 32, 32),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     2, ["f16,f16,f16,f16,f16", "f16,f16,f16,f32,f32"],
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Turing,
                     TensorOp((16, 8, 8)),
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=1),
    *gen_conv_params(ConvFwdAndBwdInput, (32, 128, 64), (32, 64, 32),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     2, ["f16,f16,f16,f16,f16", "f16,f16,f16,f32,f32"],
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Turing,
                     TensorOp((16, 8, 8)),
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=1),
    *gen_conv_params(ConvFwdAndBwdInput, (32, 128, 64), (32, 32, 64),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     2, ["f16,f16,f16,f16,f16", "f16,f16,f16,f32,f32"],
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Turing,
                     TensorOp((16, 8, 8)),
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=1),
    *gen_conv_params(ConvFwdAndBwdInput, (32, 128, 64), (32, 64, 64),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     2, ["f16,f16,f16,f16,f16", "f16,f16,f16,f32,f32"],
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Turing,
                     TensorOp((16, 8, 8)),
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=1),
    *gen_conv_params(ConvFwdAndBwdInput, (64, 128, 32), (32, 64, 32),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     2, ["f16,f16,f16,f16,f16", "f16,f16,f16,f32,f32"],
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Turing,
                     TensorOp((16, 8, 8)),
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=1),
    *gen_conv_params(ConvFwdAndBwdInput, (64, 128, 64), (32, 64, 32),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     2, ["f16,f16,f16,f16,f16", "f16,f16,f16,f32,f32"],
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Turing,
                     TensorOp((16, 8, 8)),
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=1),
    *gen_conv_params(ConvFwdAndBwdInput, (64, 64, 32), (32, 32, 32),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     2, ["f16,f16,f16,f16,f16", "f16,f16,f16,f32,f32"],
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Turing,
                     TensorOp((16, 8, 8)),
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=1),
    *gen_conv_params(ConvBwdWeight, (64, 64, 32), (32, 32, 32),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     2,
                     "f16,f16,f16,f32,f32",
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Turing,
                     TensorOp((16, 8, 8)),
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=0),
    *gen_conv_params(ConvBwdWeight, (128, 128, 32), (32, 64, 32),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     2,
                     "f16,f16,f16,f32,f32",
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Turing,
                     TensorOp((16, 8, 8)),
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=1),
    *gen_conv_params(ConvBwdWeight, (64, 64, 32), (32, 32, 32),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     2,
                     "f16,f16,f16,f32,f32",
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Turing,
                     TensorOp((16, 8, 8)),
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=1),


    *gen_conv_params(ConvBwdWeight, (64, 16, 32), (32, 16, 32),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     2,
                     "f16,f16,f16,f32,f32",
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Turing,
                     TensorOp((16, 8, 8)),
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=0),
    *gen_conv_params(ConvBwdWeight, (64, 16, 32), (32, 16, 32),
                     NDIM_DONT_CARE,
                     ConvIterAlgo.Optimized,
                     2,
                     "f16,f16,f16,f32,f32",
                     NHWC,
                     NHWC,
                     NHWC,
                     GemmAlgo.Turing,
                     TensorOp((16, 8, 8)),
                     mask_sparse=True,
                     increment_k_first=True,
                     access_per_vector=1),
    
    # *gen_conv_params(ConvBwdWeight, (32, 64, 32), (32, 32, 16), NDIM_DONT_CARE, ConvIterAlgo.Optimized, 2, "f16,f16,f16,f32,f32",
    #     NHWC, NHWC, NHWC, GemmAlgo.Turing, TensorOp((16, 8, 8)), mask_sparse=True, increment_k_first=True, access_per_vector=1),

    # gen_conv_params(ConvFwdAndBwdInput, )
]

if not SPCONV_INT8_DEBUG:
    IMPLGEMM_TURING_PARAMS.extend([ 
        *gen_conv_params(ConvFwdAndBwdInput, (64, 32, 32), (32, 32, 32),
                        NDIM_DONT_CARE,
                        ConvIterAlgo.Optimized,
                        2,
                        ["s8,s8,s8,s32,f32", "s8,s8,s8,s32,f16"],
                        NHWC,
                        NHWC,
                        NHWC,
                        GemmAlgo.Turing,
                        TensorOp((8, 8, 16)),
                        mask_sparse=True,
                        increment_k_first=True,
                        access_per_vector=1,
                        is_nvrtc=True,
                        int8_inference=True),
        *gen_conv_params(ConvFwdAndBwdInput, (32, 64, 32), (32, 32, 32),
                        NDIM_DONT_CARE,
                        ConvIterAlgo.Optimized,
                        2,
                        ["s8,s8,s8,s32,f32", "s8,s8,s8,s32,f16"],
                        NHWC,
                        NHWC,
                        NHWC,
                        GemmAlgo.Turing,
                        TensorOp((8, 8, 16)),
                        mask_sparse=True,
                        increment_k_first=True,
                        access_per_vector=1,
                        is_nvrtc=True,
                        int8_inference=True),
        *gen_conv_params(ConvFwdAndBwdInput, (32, 32, 32), (32, 32, 32),
                        NDIM_DONT_CARE,
                        ConvIterAlgo.Optimized,
                        2,
                        ["s8,s8,s8,s32,f32", "s8,s8,s8,s32,f16"],
                        NHWC,
                        NHWC,
                        NHWC,
                        GemmAlgo.Turing,
                        TensorOp((8, 8, 16)),
                        mask_sparse=True,
                        increment_k_first=True,
                        access_per_vector=1,
                        is_nvrtc=True,
                        int8_inference=True),
        *gen_conv_params(ConvFwdAndBwdInput, (32, 32, 32), (16, 32, 32),
                        NDIM_DONT_CARE,
                        ConvIterAlgo.Optimized,
                        2,
                        ["s8,s8,s8,s32,f32", "s8,s8,s8,s32,f16"],
                        NHWC,
                        NHWC,
                        NHWC,
                        GemmAlgo.Turing,
                        TensorOp((8, 8, 16)),
                        mask_sparse=True,
                        increment_k_first=True,
                        access_per_vector=1,
                        is_nvrtc=True,
                        int8_inference=True),
        *gen_conv_params(ConvFwdAndBwdInput, (32, 32, 32), (16, 16, 32),
                        NDIM_DONT_CARE,
                        ConvIterAlgo.Optimized,
                        2,
                        ["s8,s8,s8,s32,f32", "s8,s8,s8,s32,f16"],
                        NHWC,
                        NHWC,
                        NHWC,
                        GemmAlgo.Turing,
                        TensorOp((8, 8, 16)),
                        mask_sparse=True,
                        increment_k_first=True,
                        access_per_vector=1,
                        is_nvrtc=True,
                        int8_inference=True),
        *gen_conv_params(ConvFwdAndBwdInput, (32, 32, 32), (32, 16, 32),
                        NDIM_DONT_CARE,
                        ConvIterAlgo.Optimized,
                        2,
                        ["s8,s8,s8,s32,f32", "s8,s8,s8,s32,f16"],
                        NHWC,
                        NHWC,
                        NHWC,
                        GemmAlgo.Turing,
                        TensorOp((8, 8, 16)),
                        mask_sparse=True,
                        increment_k_first=True,
                        access_per_vector=1,
                        is_nvrtc=True,
                        int8_inference=True),

        *gen_conv_params(ConvFwdAndBwdInput, (64, 64, 64), (32, 32, 64),
                        NDIM_DONT_CARE,
                        ConvIterAlgo.Optimized,
                        2,
                        ["s8,s8,s8,s32,f32", "s8,s8,s8,s32,f16"],
                        NHWC,
                        NHWC,
                        NHWC,
                        GemmAlgo.Turing,
                        TensorOp((8, 8, 16)),
                        mask_sparse=True,
                        increment_k_first=True,
                        access_per_vector=1,
                        is_nvrtc=True,
                        int8_inference=True),

        *gen_conv_params(ConvFwdAndBwdInput, (64, 128, 64), (32, 64, 64),
                        NDIM_DONT_CARE,
                        ConvIterAlgo.Optimized,
                        2,
                        ["s8,s8,s8,s32,f32", "s8,s8,s8,s32,f16"],
                        NHWC,
                        NHWC,
                        NHWC,
                        GemmAlgo.Turing,
                        TensorOp((8, 8, 16)),
                        mask_sparse=True,
                        increment_k_first=True,
                        access_per_vector=1,
                        is_nvrtc=True,
                        int8_inference=True),
        
        *gen_conv_params(ConvFwdAndBwdInput, (64, 128, 32), (32, 64, 32),
                        NDIM_DONT_CARE,
                        ConvIterAlgo.Optimized,
                        2,
                        ["s8,s8,s8,s32,f32", "s8,s8,s8,s32,f16"],
                        NHWC,
                        NHWC,
                        NHWC,
                        GemmAlgo.Turing,
                        TensorOp((8, 8, 16)),
                        mask_sparse=True,
                        increment_k_first=True,
                        access_per_vector=1,
                        is_nvrtc=True,
                        int8_inference=True),

        *gen_conv_params(ConvFwdAndBwdInput, (128, 64, 64), (64, 32, 64),
                        NDIM_DONT_CARE,
                        ConvIterAlgo.Optimized,
                        2,
                        ["s8,s8,s8,s32,f32", "s8,s8,s8,s32,f16"],
                        NHWC,
                        NHWC,
                        NHWC,
                        GemmAlgo.Turing,
                        TensorOp((8, 8, 16)),
                        mask_sparse=True,
                        increment_k_first=True,
                        access_per_vector=1,
                        is_nvrtc=True,
                        int8_inference=True),
        # TODO 16,8,32 produce wrong result.
        *gen_conv_params(ConvFwdAndBwdInput, (128, 64, 32), (64, 32, 32),
                        NDIM_DONT_CARE,
                        ConvIterAlgo.Optimized,
                        2,
                        ["s8,s8,s8,s32,f32", "s8,s8,s8,s32,f16"],
                        NHWC,
                        NHWC,
                        NHWC,
                        GemmAlgo.Turing,
                        TensorOp((8, 8, 16)),
                        mask_sparse=True,
                        increment_k_first=True,
                        access_per_vector=1,
                        is_nvrtc=True,
                        int8_inference=True),

        *gen_conv_params(ConvFwdAndBwdInput, (128, 256, 64), (64, 128, 64),
                        NDIM_DONT_CARE,
                        ConvIterAlgo.Optimized,
                        2,
                        ["s8,s8,s8,s32,f32", "s8,s8,s8,s32,f16"],
                        NHWC,
                        NHWC,
                        NHWC,
                        GemmAlgo.Turing,
                        TensorOp((8, 8, 16)),
                        mask_sparse=True,
                        increment_k_first=True,
                        access_per_vector=1,
                        is_nvrtc=True,
                        int8_inference=True),
        
        *gen_conv_params(ConvFwdAndBwdInput, (256, 128, 64), (128, 64, 64),
                        NDIM_DONT_CARE,
                        ConvIterAlgo.Optimized,
                        2,
                        ["s8,s8,s8,s32,f32", "s8,s8,s8,s32,f16"],
                        NHWC,
                        NHWC,
                        NHWC,
                        GemmAlgo.Turing,
                        TensorOp((8, 8, 16)),
                        mask_sparse=True,
                        increment_k_first=True,
                        access_per_vector=1,
                        is_nvrtc=True,
                        int8_inference=True),

        *gen_conv_params(ConvFwdAndBwdInput, (128, 128, 64), (64, 64, 64),
                        NDIM_DONT_CARE,
                        ConvIterAlgo.Optimized,
                        2,
                        ["s8,s8,s8,s32,f32", "s8,s8,s8,s32,f16"],
                        NHWC,
                        NHWC,
                        NHWC,
                        GemmAlgo.Turing,
                        TensorOp((8, 8, 16)),
                        mask_sparse=True,
                        increment_k_first=True,
                        access_per_vector=1,
                        is_nvrtc=True,
                        int8_inference=True),
    ])

def _dtype_shortcuts(param):
    return tuple(
        getattr(param, name).shortcut()
        for name in ("dtype_a", "dtype_b", "dtype_c")
    )


def _all_dtypes_are(param, dtype_shortcut: str):
    return all(shortcut == dtype_shortcut for shortcut in _dtype_shortcuts(param))


def _has_any_dtype(param, dtype_shortcuts):
    return any(shortcut in dtype_shortcuts for shortcut in _dtype_shortcuts(param))


def _is_fp32_simt_param(param):
    return param.algo == GemmAlgo.Simt and _all_dtypes_are(param, "f32")


def _is_ampere_param(param):
    return param.algo == GemmAlgo.Ampere


def _is_static_param(param):
    return not getattr(param, "is_nvrtc", False)


def _is_non_int8_param(param):
    return not getattr(param, "int8_inference", False)


def _is_fp32_ampere_param(param):
    return (
        _is_ampere_param(param)
        and _is_static_param(param)
        and _is_non_int8_param(param)
        and _all_dtypes_are(param, "f32")
    )


def _is_f16_ampere_param(param):
    return (
        _is_ampere_param(param)
        and _is_static_param(param)
        and _is_non_int8_param(param)
        and _has_any_dtype(param, {"f16"})
    )


def _is_ampere_no_int8_static_param(param):
    return _is_ampere_param(param) and _is_static_param(param) and _is_non_int8_param(param)


def _is_ampere_int8_param(param):
    return _is_ampere_param(param) and getattr(param, "int8_inference", False)


def _is_non_int8_nvrtc_param(param):
    return getattr(param, "is_nvrtc", False) and _is_non_int8_param(param)


def _is_fp8_param(param):
    return _has_any_dtype(param, {"e4m3", "e5m2"})


def _filter_params(params, predicate):
    return [param for param in params if predicate(param)]


def _clear_turing_volta():
    global SHUFFLE_TURING_PARAMS, SHUFFLE_VOLTA_PARAMS
    global IMPLGEMM_TURING_PARAMS, IMPLGEMM_VOLTA_PARAMS
    SHUFFLE_TURING_PARAMS = []
    SHUFFLE_VOLTA_PARAMS = []
    IMPLGEMM_TURING_PARAMS = []
    IMPLGEMM_VOLTA_PARAMS = []


_DTK_KERNEL_FILTER = os.getenv("SPCONV_DTK_KERNEL_FILTER", "").lower()
if _DTK_KERNEL_FILTER == "dtk_smoke":
    SHUFFLE_SIMT_PARAMS = _filter_params(SHUFFLE_SIMT_PARAMS, _is_fp32_simt_param)[:4]
    SHUFFLE_AMPERE_PARAMS = []
    IMPLGEMM_SIMT_PARAMS = _filter_params(IMPLGEMM_SIMT_PARAMS, _is_fp32_simt_param)[:4]
    IMPLGEMM_AMPERE_PARAMS = []
    _clear_turing_volta()
elif _DTK_KERNEL_FILTER == "dtk_fp32_simt":
    SHUFFLE_SIMT_PARAMS = _filter_params(SHUFFLE_SIMT_PARAMS, _is_fp32_simt_param)
    SHUFFLE_AMPERE_PARAMS = []
    IMPLGEMM_SIMT_PARAMS = _filter_params(IMPLGEMM_SIMT_PARAMS, _is_fp32_simt_param)
    IMPLGEMM_AMPERE_PARAMS = []
    _clear_turing_volta()
elif _DTK_KERNEL_FILTER == "dtk_all_simt":
    SHUFFLE_AMPERE_PARAMS = []
    IMPLGEMM_AMPERE_PARAMS = []
    _clear_turing_volta()
elif _DTK_KERNEL_FILTER == "dtk_fp32_ampere":
    SHUFFLE_SIMT_PARAMS = _filter_params(SHUFFLE_SIMT_PARAMS, _is_fp32_simt_param)
    SHUFFLE_AMPERE_PARAMS = _filter_params(SHUFFLE_AMPERE_PARAMS, _is_fp32_ampere_param)
    IMPLGEMM_SIMT_PARAMS = _filter_params(IMPLGEMM_SIMT_PARAMS, _is_fp32_simt_param)
    IMPLGEMM_AMPERE_PARAMS = _filter_params(IMPLGEMM_AMPERE_PARAMS, _is_fp32_ampere_param)
    _clear_turing_volta()
elif _DTK_KERNEL_FILTER == "dtk_f16_ampere":
    SHUFFLE_SIMT_PARAMS = _filter_params(SHUFFLE_SIMT_PARAMS, _is_fp32_simt_param)
    SHUFFLE_AMPERE_PARAMS = _filter_params(SHUFFLE_AMPERE_PARAMS, _is_f16_ampere_param)
    IMPLGEMM_SIMT_PARAMS = _filter_params(IMPLGEMM_SIMT_PARAMS, _is_fp32_simt_param)
    IMPLGEMM_AMPERE_PARAMS = _filter_params(IMPLGEMM_AMPERE_PARAMS, _is_f16_ampere_param)
    _clear_turing_volta()
elif _DTK_KERNEL_FILTER == "dtk_ampere_no_int8":
    SHUFFLE_AMPERE_PARAMS = _filter_params(SHUFFLE_AMPERE_PARAMS, _is_ampere_no_int8_static_param)
    IMPLGEMM_AMPERE_PARAMS = _filter_params(IMPLGEMM_AMPERE_PARAMS, _is_ampere_no_int8_static_param)
    _clear_turing_volta()
elif _DTK_KERNEL_FILTER == "dtk_int8_ampere":
    SHUFFLE_AMPERE_PARAMS = _filter_params(
        SHUFFLE_AMPERE_PARAMS,
        lambda param: _is_ampere_no_int8_static_param(param) or _is_ampere_int8_param(param))
    IMPLGEMM_AMPERE_PARAMS = _filter_params(
        IMPLGEMM_AMPERE_PARAMS,
        lambda param: _is_ampere_no_int8_static_param(param) or _is_ampere_int8_param(param))
    _clear_turing_volta()
elif _DTK_KERNEL_FILTER == "dtk_nvrtc":
    SHUFFLE_AMPERE_PARAMS = _filter_params(
        SHUFFLE_AMPERE_PARAMS,
        lambda param: _is_ampere_no_int8_static_param(param) or _is_non_int8_nvrtc_param(param))
    IMPLGEMM_AMPERE_PARAMS = _filter_params(
        IMPLGEMM_AMPERE_PARAMS,
        lambda param: _is_ampere_no_int8_static_param(param) or _is_non_int8_nvrtc_param(param))
    _clear_turing_volta()
elif _DTK_KERNEL_FILTER == "dtk_fp8_probe":
    SHUFFLE_SIMT_PARAMS = _filter_params(SHUFFLE_SIMT_PARAMS, _is_fp32_simt_param)
    SHUFFLE_AMPERE_PARAMS = _filter_params(SHUFFLE_AMPERE_PARAMS, _is_fp8_param)
    IMPLGEMM_SIMT_PARAMS = _filter_params(IMPLGEMM_SIMT_PARAMS, _is_fp32_simt_param)
    IMPLGEMM_AMPERE_PARAMS = _filter_params(IMPLGEMM_AMPERE_PARAMS, _is_fp8_param)
    _clear_turing_volta()
elif _DTK_KERNEL_FILTER == "dtk_all_static_no_nvrtc":
    SHUFFLE_SIMT_PARAMS = _filter_params(SHUFFLE_SIMT_PARAMS, _is_static_param)
    SHUFFLE_TURING_PARAMS = _filter_params(SHUFFLE_TURING_PARAMS, _is_static_param)
    SHUFFLE_VOLTA_PARAMS = _filter_params(SHUFFLE_VOLTA_PARAMS, _is_static_param)
    SHUFFLE_AMPERE_PARAMS = _filter_params(SHUFFLE_AMPERE_PARAMS, _is_static_param)
    IMPLGEMM_SIMT_PARAMS = _filter_params(IMPLGEMM_SIMT_PARAMS, _is_static_param)
    IMPLGEMM_TURING_PARAMS = _filter_params(IMPLGEMM_TURING_PARAMS, _is_static_param)
    IMPLGEMM_VOLTA_PARAMS = _filter_params(IMPLGEMM_VOLTA_PARAMS, _is_static_param)
    IMPLGEMM_AMPERE_PARAMS = _filter_params(IMPLGEMM_AMPERE_PARAMS, _is_static_param)
elif _DTK_KERNEL_FILTER == "dtk_all":
    pass


ALL_NATIVE_PARAMS = SHUFFLE_SIMT_PARAMS + SHUFFLE_TURING_PARAMS + SHUFFLE_VOLTA_PARAMS + SHUFFLE_AMPERE_PARAMS

ALL_IMPGEMM_PARAMS = IMPLGEMM_SIMT_PARAMS + IMPLGEMM_TURING_PARAMS + IMPLGEMM_VOLTA_PARAMS + IMPLGEMM_AMPERE_PARAMS
