# 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.

import enum
from functools import reduce
from itertools import product

import numpy as np
import pytest
from tvm_ffi import Shape

import tvm
import tvm.testing
from tvm.runtime import DataType, disco

if disco is None:
    pytest.skip("disco runtime is not available", allow_module_level=True)


class AllReduceStrategyType(enum.IntEnum):
    RING = 0
    ONESHOT = 1
    TWOSHOT = 2
    AUTO = 3


_shapes = [(2, 3), (3, 4), (128, 128)]

_strategies = [
    AllReduceStrategyType.RING,
    AllReduceStrategyType.ONESHOT,
    AllReduceStrategyType.TWOSHOT,
    AllReduceStrategyType.AUTO,
]

_compiled_ccl = tvm.get_global_func("runtime.disco.compiled_ccl", allow_missing=True)
if _compiled_ccl is None:
    pytest.skip("Disco CCL is not enabled in this TVM build", allow_module_level=True)
_ccl = [ccl for ccl in _compiled_ccl() if ccl == "nccl"]


@pytest.mark.parametrize("shape", _shapes)
@pytest.mark.parametrize("ccl", _ccl)
@pytest.mark.parametrize("strategy", _strategies)
def test_allreduce(shape, ccl, strategy):
    devices = [0, 1]
    sess = disco.ProcessSession(num_workers=len(devices))
    sess.init_ccl(ccl, *devices)

    num_elements = reduce(lambda x, y: x * y, shape)
    dtype = "float32"
    falloc_ipc_storage = sess.get_global_func("runtime.disco.cuda_ipc.alloc_storage")
    falloc_tensor = sess.get_global_func("vm.builtin.alloc_tensor")
    fallreduce = sess.get_global_func("runtime.disco.cuda_ipc.custom_allreduce")
    d_storage = sess.call_packed(falloc_ipc_storage, Shape(shape), DataType(dtype))
    d_input = sess.call_packed(falloc_tensor, d_storage, 0, Shape(shape), DataType(dtype))

    array_1 = np.arange(num_elements, dtype="float32").reshape(*shape)
    array_2 = np.arange(start=1, stop=-(num_elements - 1), step=-1, dtype="float32").reshape(*shape)
    d_input.debug_copy_from(0, array_1)
    d_input.debug_copy_from(1, array_2)
    d_output = sess.empty(shape, "float32")

    sess.call_packed(fallreduce, d_input, strategy, d_output)
    result_1 = d_output.debug_get_from_remote(0).numpy()
    result_2 = d_output.debug_get_from_remote(1).numpy()
    expected = np.add(array_1, array_2)
    np.testing.assert_equal(result_1, expected)
    np.testing.assert_equal(result_2, expected)


if __name__ == "__main__":
    for shape, strategy in product(_shapes, _strategies):
        test_allreduce(shape, "nccl", strategy)