# Copyright 2023 The IREE Authors
#
# Licensed under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

if(NOT IREE_BUILD_COMPILER OR NOT IREE_ENABLE_COLLECTIVE_RUNTIME_TESTS)
  return()
endif()

if(IREE_TARGET_BACKEND_CUDA AND IREE_HAL_DRIVER_CUDA)

  set(COMMON_ARGS
    "--target_backend=cuda"
    "--driver=cuda"
    "--iree_compiler_args=--iree-cuda-target=sm_53"
  )

  set(COMMON_LABELS
    "requires-gpu-nvidia"
    "driver=cuda"
  )

  iree_py_test(
    NAME
      collectives_test_cuda_1_gpu
    SRCS
      "collectives_test.py"
    ARGS
      "-k" "SingleRank"
      ${COMMON_ARGS}
    LABELS
      ${COMMON_LABELS}
  )

  iree_py_test(
    NAME
      collectives_test_cuda_2_gpus
    SRCS
      "collectives_test.py"
    ARGS
      "-k" "TwoRanks"
      ${COMMON_ARGS}
    LABELS
      ${COMMON_LABELS}
      # The NCCL collectives backend requires 1 GPU per rank.
      # To properly test for 2 ranks we need 2 GPUs.
      "requires-multiple-devices"
  )

  iree_py_test(
    NAME
      collectives_test_cuda_4_gpus
    SRCS
      "collectives_test.py"
    ARGS
      "-k" "FourRanks"
      ${COMMON_ARGS}
    LABELS
      ${COMMON_LABELS}
      "requires-multiple-devices"
  )
endif()

if(IREE_TARGET_BACKEND_ROCM AND IREE_HAL_DRIVER_HIP AND IREE_HIP_TEST_TARGET_CHIP)
  set(COMMON_ARGS
    "--target_backend=rocm"
    "--driver=hip"
    "--iree_compiler_args=--iree-hip-target=${IREE_HIP_TEST_TARGET_CHIP}"
  )

  set(COMMON_LABELS
    "requires-gpu-amd"
    "driver=hip"
  )

  iree_py_test(
    NAME
      collectives_test_hip_1_gpu
    SRCS
      "collectives_test.py"
    ARGS
      "-k" "SingleRank"
      ${COMMON_ARGS}
    LABELS
      ${COMMON_LABELS}
  )

  iree_py_test(
    NAME
      collectives_test_hip_2_gpus
    SRCS
      "collectives_test.py"
    ARGS
      "-k" "TwoRanks"
      ${COMMON_ARGS}
    LABELS
      ${COMMON_LABELS}
      # The NCCL collectives backend requires 1 GPU per rank.
      # To properly test for 2 ranks we need 2 GPUs.
      "requires-multiple-devices"
  )

  iree_py_test(
    NAME
      collectives_test_hip_4_gpus
    SRCS
      "collectives_test.py"
    ARGS
      "-k" "FourRanks"
      ${COMMON_ARGS}
    LABELS
      ${COMMON_LABELS}
      "requires-multiple-devices"
  )
endif()