# IT4Innovations
# LK 2023

name = 'PyTorch'
version = '1.10.0'

homepage = 'https://pytorch.org/'
description = """Tensors and Dynamic neural networks in Python with strong GPU acceleration.
PyTorch is a deep learning framework that puts Python first."""

toolchain = {'name': 'fosscuda', 'version': '2020b'}

source_urls = [GITHUB_RELEASE]
sources = ['%(namelower)s-v%(version)s.tar.gz']

patches = [
    'PyTorch-1.7.0_avoid-nan-in-test-torch.patch',
    'PyTorch-1.7.0_disable-dev-shm-test.patch',
    'PyTorch-1.7.1_correctly-pass-jit_opt_level.patch',
    'PyTorch-1.8.1_dont-use-gpu-ccc-in-test.patch',
    'PyTorch-1.8.1_increase-distributed-test-timeout.patch',
    'PyTorch-1.9.0_limit-world-size-for-zero-redundancy-opt-test.patch',
    'PyTorch-1.10.0_add-dummy-bfloat16-for-ppc.patch',
    'PyTorch-1.10.0_fix-alias-violation-in-bitwise-ops.patch',
    'PyTorch-1.10.0_fix-faulty-asserts.patch',
    'PyTorch-1.10.0_fix-fp16-quantization-without-fbgemm.patch',
    'PyTorch-1.10.0_fix-kineto-crash.patch',
    'PyTorch-1.10.0_fix-missing-VSX-dispatch.patch',
    'PyTorch-1.10.0_fix-test-cond-cpu.patch',
    'PyTorch-1.10.0_fix-test-dataloader-fixed-affinity.patch',
    'PyTorch-1.10.0_fix-test-model_dump.patch',
    'PyTorch-1.10.0_fix-vnni-detection.patch',
    'PyTorch-1.10.0_fix-vsx-vector-functions.patch',
    'PyTorch-1.10.0_fix-XNNPACK-tests.patch',
    'PyTorch-1.10.0_increase_zero_optimizer_test_tolerance.patch',
    'PyTorch-1.10.0_remove-vec-dump-functions.patch',
    'PyTorch-1.10.0_skip_cmake_rpath.patch',
    'PyTorch-1.10.0_skip_failing_ops_tests.patch',
    'PyTorch-1.10.0_skip_nan_tests_openblas.patch',
    'PyTorch-1.10.0_skip-nnapi-test-without-qnnpack.patch',
]
checksums = [
    '1970225700d1f9791eb553855ad3c22e7c9fd69cff2765d532506b33e449831e',  # pytorch-v1.10.0.tar.gz
    'b899aa94d9e60f11ee75a706563312ccefa9cf432756c470caa8e623991c8f18',  # PyTorch-1.7.0_avoid-nan-in-test-torch.patch
    '622cb1eaeadc06e13128a862d9946bcc1f1edd3d02b259c56a9aecc4d5406b8a',  # PyTorch-1.7.0_disable-dev-shm-test.patch
    # PyTorch-1.7.1_correctly-pass-jit_opt_level.patch
    'd4d967d47f8a6172fcbf57f0a61835482968850967c4fdb01108b720696a988d',
    '89ac7a8e9e7df2e64cf8404fe3a279f5e9b759fee41c9de3aaff9c22f385c2c6',  # PyTorch-1.8.1_dont-use-gpu-ccc-in-test.patch
    # PyTorch-1.8.1_increase-distributed-test-timeout.patch
    '7a6e512274f0b8673f4f207a5bc53387d88be7e79833f42d20365668b2118071',
    # PyTorch-1.9.0_limit-world-size-for-zero-redundancy-opt-test.patch
    'ff573660913ce055e24cfd194ce747ba5685091c631cfd443eae2a99d56b57ea',
    # PyTorch-1.10.0_add-dummy-bfloat16-for-ppc.patch
    '0c8bf4b284c5c2cb4d175645d44bed674cca5218de741a8ea698cef2884a86da',
    # PyTorch-1.10.0_fix-alias-violation-in-bitwise-ops.patch
    '426c9ead1a74b656748d4c8bf8afd4303d8b9f2394ad22b21a845d07c8ca1d12',
    # PyTorch-1.10.0_fix-faulty-asserts.patch
    '44d2bb998ec3b3b65659f3327894745888a58f957a191cb524b4003018db51a2',
    # PyTorch-1.10.0_fix-fp16-quantization-without-fbgemm.patch
    'dfc09c2e787e757191a0ebf68ace46e2a90d71263ba84e2ee764b596d2ee6769',
    # PyTorch-1.10.0_fix-kineto-crash.patch
    'dc467333b28162149af8f675929d8c6bf219f23230bfc0d39af02ba4f6f882eb',
    # PyTorch-1.10.0_fix-missing-VSX-dispatch.patch
    '29368e69e98cb16659a2043b7a0df0411c99bc57a70405f8b5300f68c25f3ee9',
    # PyTorch-1.10.0_fix-test-cond-cpu.patch
    '51f83f5d5ef69656ef35b73f17e0671e70113798421be11ea4c7b56ffcc4da03',
    # PyTorch-1.10.0_fix-test-dataloader-fixed-affinity.patch
    '313dca681f45ce3bc7c4557fdcdcbe0b77216d2c708fa30a2ec0e22c44876707',
    # PyTorch-1.10.0_fix-test-model_dump.patch
    '339148ae1a028cda6e750ac93fa38a599f66c7abe26586c9219f1a206ea14557',
    # PyTorch-1.10.0_fix-vnni-detection.patch
    '1f3664c0febfa2a3fc4c0cd3bae185f289716ac0b6c3d7e8fa1cee19ba62b7cc',
    # PyTorch-1.10.0_fix-vsx-vector-functions.patch
    '7bef5f96cb83b2d655d2f76dd7468a171d446f0b3e06da2232ec7f886484d312',
    # PyTorch-1.10.0_fix-XNNPACK-tests.patch
    'd3e749a2a42efce463e3b8a1aebb21f0edf2256682c4417297d9a44a6210e5f8',
    # PyTorch-1.10.0_increase_zero_optimizer_test_tolerance.patch
    'e65afb01786f7f030ccb5faada1eb474bb0c418bcadcf1baaa71a4fa2f3f4240',
    # PyTorch-1.10.0_remove-vec-dump-functions.patch
    '5e796c83f74e77aa193fef2720a7a41e3e0eb39df726a6f0a2a1c55d37d8fcc8',
    # PyTorch-1.10.0_skip_cmake_rpath.patch
    'ac05943bb205623f91ef140aa00869efc5fe844184bd666bebf5405808610448',
    # PyTorch-1.10.0_skip_failing_ops_tests.patch
    '399af94ffcef4a6db5226552c46f11e9b0f0f371b2d7924b9e5764d2281581ab',
    # PyTorch-1.10.0_skip_nan_tests_openblas.patch
    '7d3f83e3056d9e47a460790313238f28708beb596cafaa7ae55e374d368bbedf',
    # PyTorch-1.10.0_skip-nnapi-test-without-qnnpack.patch
    '34ba476a7bcddec323bf9eca083cb4623d0f569d081aa3add3769c24f22849d2',
]

osdependencies = [OS_PKG_IBVERBS_DEV]

builddependencies = [
    ('CMake', '3.20.1'),   # Needs 3.20 or newer.
    ('hypothesis', '5.41.5'),
]

dependencies = [
    ('Ninja', '1.10.1'),  # Required for JIT compilation of C++ extensions
    ('Python', '3.8.6'),
    ('protobuf', '3.14.0'),
    ('protobuf-python', '3.14.0'),
    ('pybind11', '2.6.0'),
 #  ('SciPy-bundle', '2020.11'),
    ('typing-extensions', '3.7.4.3'),
    ('PyYAML', '5.3.1'),
    ('MPFR', '4.1.0'),
    ('GMP', '6.2.0'),
    ('numactl', '2.0.13'),
    ('FFmpeg', '4.3.1'),
    ('Pillow', '8.0.1'),
    ('expecttest', '0.1.3'),
    ('cuDNN', '8.0.4.30', '-CUDA-%(cudaver)s', SYSTEM),
    ('magma', '2.5.4'),
    ('NCCL', '2.8.3', '-CUDA-%(cudaver)s'),
]

# default CUDA compute capabilities to use (override via --cuda-compute-capabilities)
cuda_compute_capabilities = ['3.5', '3.7', '5.2', '6.0', '6.1', '7.0', '7.2', '7.5', '8.0', '8.6']

custom_opts = ["USE_CUPTI_SO=1"]

excluded_tests = {
    '': [
        # Bad tests: https://github.com/pytorch/pytorch/issues/60260
        'distributed/elastic/utils/distributed_test',
        'distributed/elastic/multiprocessing/api_test',
        # These tests fail on A10s at the very least, they time out forever no matter how long the timeout is.
        # Possibly related to NCCL 2.8.3: https://docs.nvidia.com/deeplearning/nccl/release-notes/rel_2-8-3.html
        # 'distributed/test_distributed_fork',
        'distributed/test_distributed_spawn',
        # Fails on A10s: https://github.com/pytorch/pytorch/issues/63079
        'test_optim',
        # Test from this suite timeout often. The process group backend is deprecated anyway
        # 'distributed/rpc/test_process_group_agent',
    ]
}

runtest = 'cd test && PYTHONUNBUFFERED=1 %(python)s run_test.py --continue-through-error  --verbose %(excluded_tests)s'

# The readelf sanity check can be taken out once the TestRPATH test from https://github.com/pytorch/pytorch/pull/68912
# is accepted, since it is then checked as part of the PyTorch test suite
local_libcaffe2 = "$EBROOTPYTORCH/lib/python%%(pyshortver)s/site-packages/torch/lib/libcaffe2_nvrtc.%s" % SHLIB_EXT
sanity_check_commands = [
    "python -c 'import caffe2.python'",
    "readelf -d %s | egrep 'RPATH|RUNPATH' | grep -v stubs" % local_libcaffe2,
]
tests = ['PyTorch-check-cpp-extension.py']

moduleclass = 'devel'