# -*- Python -*-
import os
import platform
import re
import subprocess
import struct

import lit.formats
import lit.util

# Configuration file for the 'lit' test runner.
site_cfg = lit_config.params.get('site_config', None)
lit_config.load_config(config, site_cfg)

config.excludes = ['cmdparser.hpp']
config.excludes.append('spatial_batch_norm_op.h')
config.excludes.append('common_cudnn.h')
config.excludes.append('inc.h')
config.excludes.append('inet.h')
config.excludes.append('types.h')
config.excludes.append('socket.h')

delimiter = "===============================================================";
print(delimiter)
print("CUDA " + config.cuda_version_full + " - will be used for testing")
print("LLVM " + config.llvm_version + " - will be used for testing")
print(platform.machine() + " - Platform architecture")
print(platform.system() + " " + platform.release() + " - Platform OS")
print(str(config.pointer_size * 8) + " - hipify-clang binary bitness")
print(str(struct.calcsize("P") * 8) + " - python " + str(platform.python_version()) + " binary bitness")
print(delimiter)

warns = None

if not config.cuda_dnn_root or config.cuda_dnn_root == "OFF":
    config.excludes.append('cudnn_convolution_forward.cu')
    config.excludes.append('cudnn_softmax.cu')
    config.excludes.append('cudnn2miopen.cu')
    config.excludes.append('cudnn2miopen_before_9000.cu')
    print("WARN: cuDNN tests are excluded because CUDA_DNN_ROOT_DIR is not specified")
    warns = True
else:
    if config.cudnn_version_major >= 9:
        config.excludes.append('cudnn2miopen_before_9000.cu')

if not config.cuda_cub_root or config.cuda_cub_root == "OFF":
    config.excludes.append('cub_01.cu')
    config.excludes.append('cub_02.cu')
    config.excludes.append('cub_03.cu')
    print("WARN: CUB tests are excluded because CUDA_CUB_ROOT_DIR is not specified")
    warns = True
else:
    if config.cuda_version_major == 12 and config.cuda_version_minor >= 6:
        config.excludes.append('cub_01.cu')
        config.excludes.append('cub_02.cu')
        config.excludes.append('cub_03.cu')

if not config.cuda_tensor_root or config.cuda_tensor_root == "OFF":
    config.excludes.append('cutensor2hiptensor.cu')
    config.excludes.append('cutensor2hiptensor_before_20000.cu')
    print("WARN: cuTensor tests are excluded because CUDA_TENSOR_ROOT_DIR is not specified")
    warns = True
else:
    if config.cutensor_version_major >= 2:
       config.excludes.append('cutensor2hiptensor_before_20000.cu')

if not config.cuda_file_root or config.cuda_file_root == "OFF" or config.cuda_version_major < 11 or (config.cuda_version_major == 11 and config.cuda_version_minor < 4):
    config.excludes.append('cufile2hipfile.cu')
    config.excludes.append('cufile2hipfile_functions.cu')
    print("WARN: cuFile tests are excluded because CUDA_FILE_ROOT_DIR is not specified or CUDA version < 11.4")
    warns = True

if warns:
    print(delimiter)

clang_arguments = "-v"

if config.cuda_version_major == 7 and config.cuda_version_minor == 0:
    config.excludes.append('headers_test_09.cu')
    config.excludes.append('cudnn_convolution_forward.cu')
    config.excludes.append('cusparse2rocsparse_7050.cu')
    config.excludes.append('headers_test_12_SOLVER_7050.cu')
    config.excludes.append('headers_test_09_fp16_7050.cu')

if config.cuda_version_major < 8:
    config.excludes.append('cuSPARSE_02.cu')

if config.cuda_version_major < 9:
    config.excludes.append('cuSPARSE_04.cu')
    config.excludes.append('cuSPARSE_05.cu')
    config.excludes.append('cuSPARSE_06.cu')
    config.excludes.append('cuSPARSE_07.cu')
    config.excludes.append('benchmark_curand_kernel.cpp')
    config.excludes.append('reinterpret_cast.cu')
    config.excludes.append('half2_allocators.cu')
    config.excludes.append('cub_01.cu')
    config.excludes.append('cub_02.cu')
    config.excludes.append('cub_03.cu')
    config.excludes.append('runtime_functions_9000.cu')

if config.cuda_version_major < 9 or (config.cuda_version_major == 9 and config.cuda_version_minor < 2):
    config.excludes.append('cusparse2rocsparse_9020.cu')
    config.excludes.append('cusparse2rocsparse_9020_12000.cu')

if config.cuda_version_major < 10:
    config.excludes.append('cuSPARSE_08.cu')
    config.excludes.append('cuSPARSE_09.cu')
    config.excludes.append('cuSPARSE_10.cu')
    config.excludes.append('cuSPARSE_11.cu')
    config.excludes.append('simple_mechs.cu')
    config.excludes.append('cusparse2rocsparse_10000.cu')

if config.cuda_version_major < 10 or (config.cuda_version_major == 10 and config.cuda_version_minor < 1):
    config.excludes.append('cusparse2rocsparse_10010.cu')

if config.cuda_version_major < 10 or (config.cuda_version_major == 10 and config.cuda_version_minor < 1) or config.cuda_version_major >= 12:
    config.excludes.append('cusparse2rocsparse_10010_12000.cu')

if config.cuda_version_major < 10 or (config.cuda_version_major == 10 and config.cuda_version_minor < 1):
    config.excludes.append('cublaslt2hipblaslt.cu')
    config.excludes.append('cublaslt2hipblaslt_10010_10020.cu')

if not (config.cuda_version_major == 10 and (config.cuda_version_minor == 1 and config.cuda_version_minor == 2)):
    config.excludes.append('cublaslt2hipblaslt_10010_10020.cu')

if config.cuda_version_major < 11 and sys.platform in ['win32']:
    config.excludes.append('cusparse2rocsparse_10010_12000.cu')

# [NOTE] cmake doesn't load lib VERSIONs, including CUSPARSE_VERSION, which equals 11200 for CUDA 11.1.0 and 11300 for CUDA 11.1.1
# [NOTE] Both CUDA 11.1.0 CUDA 11.1.1 have the same CUDA_VERSION 11010, and we can't distinguish them, thus exclude the below tests from both 11.1.0 and 11.1.1
if config.cuda_version_major < 11 or (config.cuda_version_major == 11 and config.cuda_version_minor <= 1) or config.cuda_version_major >= 12:
    config.excludes.append('cusparse2rocsparse_11010_12000.cu')

# [NOTE] cmake doesn't load lib VERSIONs, including CUSPARSE_VERSION, which equals 11500 for CUDA 11.3.0 and 11600 for CUDA 11.3.1
# [NOTE] Both CUDA 11.3.0 CUDA 11.3.1 have the same CUDA_VERSION 11030, and we can't distinguish them, thus exclude the below tests from both 11.3.0 and 11.3.1
if config.cuda_version_major < 11 or (config.cuda_version_major == 11 and config.cuda_version_minor <= 3) or config.cuda_version_major >= 12:
    config.excludes.append('cusparse2rocsparse_11030_12000.cu')
    config.excludes.append('cusparse2hipsparse_11030_12000.cu')

if config.cuda_version_major <= 10:
    config.excludes.append('headers_test_12_SOLVER_10010.cu')

if config.cuda_version_major > 10:
    clang_arguments += " -DTHRUST_IGNORE_CUB_VERSION_CHECK"
    config.excludes.append('cudnn_convolution_forward.cu')
    config.excludes.append('cuSPARSE_01.cu')
    config.excludes.append('cuSPARSE_02.cu')
    config.excludes.append('cuSPARSE_12.cu')

if config.cuda_version_major < 11:
    config.excludes.append('headers_test_09_bf16_11000.cu')

if config.cuda_version_major < 11 or (config.cuda_version_major == 11 and config.cuda_version_minor < 1):
    config.excludes.append('runtime_functions_11010.cu')

if config.cuda_version_major >= 11:
    config.excludes.append('cusparse2rocsparse_before_11000.cu')

if config.cuda_version_major < 11 or (config.cuda_version_major == 11 and config.cuda_version_minor < 8):
    config.excludes.append('cudevice2hipdevice.cu')
    config.excludes.append('cudevice2hipdevice_before_11080_after_12011.cu')

if (config.cuda_version_major == 11 and config.cuda_version_minor >=8) or (config.cuda_version_major == 12 and config.cuda_version_minor < 2):
    config.excludes.append('cudevice2hipdevice_before_11080_after_12011.cu')

if (config.cuda_version_major < 12) or (config.cuda_version_major == 12 and config.cuda_version_minor < 8):
    config.excludes.append('cudevice2hipdevice_12080.cu')
    config.excludes.append('runtime_functions_12080.cu')

if config.cuda_version_major < 12:
    config.excludes.append('headers_test_06_12000.cu')
    config.excludes.append('headers_test_07_12000.cu')
    config.excludes.append('headers_test_08_12000.cu')
    config.excludes.append('headers_test_09_12000.cu')
    config.excludes.append('runtime_functions_12000.cu')
    config.excludes.append('cusparse2rocsparse_12000.cu')
    config.excludes.append('cusparse2hipsparse_12000.cu')

if config.cuda_version_major >= 12:
    config.excludes.append('headers_test_06.cu')
    config.excludes.append('headers_test_07.cu')
    config.excludes.append('headers_test_08.cu')
    config.excludes.append('headers_test_09.cu')
    config.excludes.append('tex2dKernel.cpp')
    config.excludes.append('cuSPARSE_11.cu')
    config.excludes.append('cublas_0_based_indexing.cu')
    config.excludes.append('cublas_0_based_indexing_rocblas.cu')
    config.excludes.append('cusparse2rocsparse_9020_12000.cu')
    config.excludes.append('cusparse2rocsparse_before_12000.cu')
    clang_arguments += " -D_LIBCUDACXX_OBJECT_FORMAT_COFF"

if config.cuda_version_major >= 13:
    config.excludes.append('cd_intro_before_13000.cu')
    config.excludes.append('driver_functions_before_13000.cu')
    config.excludes.append('runtime_functions_before_13000.cu')
    config.excludes.append('intro_before_13000.cu')
    config.excludes.append('headers_test_09_rocrand_before_13000.cu')

if config.cuda_version_major < 13:
    config.excludes.append('cd_intro.cu')
    config.excludes.append('intro.cu')

if config.llvm_version_major < 8:
    config.excludes.append('cd_intro.cu')
    config.excludes.append('cd_intro_before_13000.cu')

if config.llvm_version_major < 10:
    config.excludes.append('pp_if_else_conditionals_LLVM_10.cu')
    config.excludes.append('pp_if_else_conditionals_01_LLVM_10.cu')

if config.llvm_version_major > 18:
    clang_arguments += " -Wno-missing-template-arg-list-after-template-kw"

# name: The name of this test suite.
config.name = 'hipify'

# suffixes: CUDA source is only supported
config.suffixes = ['.cu','.cuh','.cpp','.c','.hpp','.h']

# testFormat: The test format to use to interpret tests.
config.test_format = lit.formats.ShTest()

# test_source_root: The root path where tests are located.
config.test_source_root = os.path.dirname(__file__)

# test_exec_root: The path where tests are located (default is the test suite root).
#config.test_exec_root = config.test_source_root

# target_triple: Used by ShTest and TclTest formats for XFAIL checks.
config.target_triple = '(unused)'

# available_features: Used by ShTest and TclTest formats for REQUIRES checks.
config.available_features = []

obj_root = getattr(config, 'obj_root', None)
if obj_root is not None:
    config.test_exec_root = obj_root

if obj_root is not None:
    llvm_tools_dir = getattr(config, 'llvm_tools_dir', None)
    if not llvm_tools_dir:
        lit_config.fatal('No LLVM tools dir set!')
    path = os.path.pathsep.join((llvm_tools_dir, config.environment['PATH']))
    config.environment['PATH'] = path

if config.hipify_clang_tests_only and config.hipify_clang_tests_only != "0" and config.hipify_clang_tests_only.upper() != "OFF":
    hipify_path = config.hipify_install_path
else:
    hipify_path = obj_root

inc_subfolder = ""
if config.cuda_version_major < 11 or (config.cuda_version_major == 11 and config.cuda_version_minor < 6):
    inc_subfolder = "/inc"

if sys.platform in ['win32']:
    run_test_ext = ".bat"
    if not config.hipify_clang_tests_only or config.hipify_clang_tests_only == "0" or config.hipify_clang_tests_only.upper() == "OFF":
        hipify_path += "/" + config.build_type
else:
    run_test_ext = ".sh"

if config.pointer_size == 8:
    clang_arguments += " -D__LP64__"

# Set max clang's CudaArch for corresponding CUDA version
# to support maximum CUDA features in offline tests
if config.cuda_version_major == 7:
    if config.cuda_version_minor == 5:
        clang_arguments += " --cuda-gpu-arch=sm_53"
    else:
        clang_arguments += " --cuda-gpu-arch=sm_52"
elif config.cuda_version_major == 8:
    clang_arguments += " --cuda-gpu-arch=sm_62"
elif config.cuda_version_major == 9:
    if config.cuda_version_minor == 2 and config.llvm_version_major >= 7 and config.llvm_version_major < 8:
        clang_arguments += " --cuda-gpu-arch=sm_72"
    elif config.llvm_version_major >= 6 and config.llvm_version_major < 7:
        clang_arguments += " --cuda-gpu-arch=sm_70"
elif config.cuda_version_major == 10:
    if config.llvm_version_major >= 8 and config.llvm_version_major < 11:
        clang_arguments += " --cuda-gpu-arch=sm_75"
elif config.cuda_version_major == 11:
    if config.cuda_version_minor == 0 and config.llvm_version_major >= 11 and config.llvm_version_major < 13:
        clang_arguments += " --cuda-gpu-arch=sm_80"
    if config.cuda_version_minor > 0 and config.llvm_version_major >= 13:
        clang_arguments += " --cuda-gpu-arch=sm_86"

# CUDA SDK ROOT
if config.cuda_sdk_root and config.cuda_sdk_root != "OFF":
    clang_arguments += " -isystem'%s'/common" + inc_subfolder

# cuDNN ROOT
if config.cuda_dnn_root and config.cuda_dnn_root != "OFF":
    clang_arguments += " -I'%s'"

# CUB ROOT
if config.cuda_cub_root and config.cuda_cub_root != "OFF":
    clang_arguments += " -I'%s'"

# cuTensor ROOT
if config.cuda_tensor_root and config.cuda_tensor_root != "OFF":
    clang_arguments += " -I'%s'"

# cuFile ROOT
if config.cuda_file_root and config.cuda_file_root != "OFF":
    clang_arguments += " -I'%s'"

i_subst = 0
if config.cuda_sdk_root != "OFF":
    i_subst = 1
if config.cuda_dnn_root != "OFF":
    i_subst += 2
if config.cuda_cub_root != "OFF":
    i_subst += 4
if config.cuda_tensor_root != "OFF":
    i_subst += 8
if config.cuda_file_root != "OFF":
    i_subst += 16

if i_subst == 31:
    config.substitutions.append(("%clang_args", clang_arguments % (config.cuda_sdk_root, config.cuda_dnn_root, config.cuda_cub_root, config.cuda_tensor_root, config.cuda_file_root)))
elif i_subst == 30:
    config.substitutions.append(("%clang_args", clang_arguments % (config.cuda_dnn_root, config.cuda_cub_root, config.cuda_tensor_root, config.cuda_file_root)))
elif i_subst == 29:
    config.substitutions.append(("%clang_args", clang_arguments % (config.cuda_sdk_root, config.cuda_cub_root, config.cuda_tensor_root, config.cuda_file_root)))
elif i_subst == 28:
    config.substitutions.append(("%clang_args", clang_arguments % (config.cuda_cub_root, config.cuda_tensor_root, config.cuda_file_root)))
elif i_subst == 27:
    config.substitutions.append(("%clang_args", clang_arguments % (config.cuda_sdk_root, config.cuda_dnn_root, config.cuda_tensor_root, config.cuda_file_root)))
elif i_subst == 26:
    config.substitutions.append(("%clang_args", clang_arguments % (config.cuda_dnn_root, config.cuda_tensor_root, config.cuda_file_root)))
elif i_subst == 25:
    config.substitutions.append(("%clang_args", clang_arguments % (config.cuda_sdk_root, config.cuda_tensor_root, config.cuda_file_root)))
elif i_subst == 24:
    config.substitutions.append(("%clang_args", clang_arguments % (config.cuda_tensor_root, config.cuda_file_root)))
elif i_subst == 23:
    config.substitutions.append(("%clang_args", clang_arguments % (config.cuda_sdk_root, config.cuda_dnn_root, config.cuda_cub_root, config.cuda_file_root)))
elif i_subst == 22:
    config.substitutions.append(("%clang_args", clang_arguments % (config.cuda_dnn_root, config.cuda_cub_root, config.cuda_file_root)))
elif i_subst == 21:
    config.substitutions.append(("%clang_args", clang_arguments % (config.cuda_sdk_root, config.cuda_cub_root, config.cuda_file_root)))
elif i_subst == 20:
    config.substitutions.append(("%clang_args", clang_arguments % (config.cuda_cub_root, config.cuda_file_root)))
elif i_subst == 19:
    config.substitutions.append(("%clang_args", clang_arguments % (config.cuda_sdk_root, config.cuda_dnn_root, config.cuda_file_root)))
elif i_subst == 18:
    config.substitutions.append(("%clang_args", clang_arguments % (config.cuda_dnn_root, config.cuda_file_root)))
elif i_subst == 17:
    config.substitutions.append(("%clang_args", clang_arguments % (config.cuda_sdk_root, config.cuda_file_root)))
elif i_subst == 16:
    config.substitutions.append(("%clang_args", clang_arguments % config.cuda_file_root))
elif i_subst == 15:
    config.substitutions.append(("%clang_args", clang_arguments % (config.cuda_sdk_root, config.cuda_dnn_root, config.cuda_cub_root, config.cuda_tensor_root)))
elif i_subst == 14:
    config.substitutions.append(("%clang_args", clang_arguments % (config.cuda_dnn_root, config.cuda_cub_root, config.cuda_tensor_root)))
elif i_subst == 13:
    config.substitutions.append(("%clang_args", clang_arguments % (config.cuda_sdk_root, config.cuda_cub_root, config.cuda_tensor_root)))
elif i_subst == 12:
    config.substitutions.append(("%clang_args", clang_arguments % (config.cuda_cub_root, config.cuda_tensor_root)))
elif i_subst == 11:
    config.substitutions.append(("%clang_args", clang_arguments % (config.cuda_sdk_root, config.cuda_dnn_root, config.cuda_tensor_root)))
elif i_subst == 10:
    config.substitutions.append(("%clang_args", clang_arguments % (config.cuda_dnn_root, config.cuda_tensor_root)))
elif i_subst == 9:
    config.substitutions.append(("%clang_args", clang_arguments % (config.cuda_sdk_root, config.cuda_tensor_root)))
elif i_subst == 8:
    config.substitutions.append(("%clang_args", clang_arguments % config.cuda_tensor_root))
elif i_subst == 7:
    config.substitutions.append(("%clang_args", clang_arguments % (config.cuda_sdk_root, config.cuda_dnn_root, config.cuda_cub_root)))
elif i_subst == 6:
    config.substitutions.append(("%clang_args", clang_arguments % (config.cuda_dnn_root, config.cuda_cub_root)))
elif i_subst == 5:
    config.substitutions.append(("%clang_args", clang_arguments % (config.cuda_sdk_root, config.cuda_cub_root)))
elif i_subst == 4:
    config.substitutions.append(("%clang_args", clang_arguments % config.cuda_cub_root))
elif i_subst == 3:
    config.substitutions.append(("%clang_args", clang_arguments % (config.cuda_sdk_root, config.cuda_dnn_root)))
elif i_subst == 2:
    config.substitutions.append(("%clang_args", clang_arguments % config.cuda_dnn_root))
elif i_subst == 1:
    config.substitutions.append(("%clang_args", clang_arguments % config.cuda_sdk_root))

if config.llvm_version_major < 4:
    hipify_arguments = "-I'%s'/include"
else:
    hipify_arguments = "\"--cuda-path=%s\""

hipify_arguments += " \"%s\""

if config.cuda_root and config.cuda_root != "OFF":
    config.substitutions.append(("%hipify_args", hipify_arguments % (config.cuda_root, config.clang_resource_dir)))
config.substitutions.append(("hipify", '"' + hipify_path + "/hipify-clang" + '"'))
config.substitutions.append(("%run_test", '"' + config.test_source_root + "/run_test" + run_test_ext + '"'))