/*
    Copyright Intel Corporation.
    
    This software and the related documents are Intel copyrighted materials, and
    your use of them is governed by the express license under which they were
    provided to you (License). Unless the License provides otherwise, you may
    not use, modify, copy, publish, distribute, disclose or transmit this
    software or the related documents without Intel's prior written permission.
    
    This software and the related documents are provided as is, with no express
    or implied warranties, other than those that are expressly stated in the
    License.
*/
#pragma once

#include "cpu_coll.hpp"
#include "reduce_scatter_strategy.hpp"

struct reduce_scatter_strategy_impl {
    static constexpr const char* class_name() {
        return "reduce_scatter";
    }

    size_t get_send_multiplier() {
        return 1;
    }

    size_t get_recv_multiplier() {
        return 1;
    }

    static const ccl::reduce_scatter_attr& get_op_attr(const bench_exec_attr& bench_attr) {
        return bench_attr.get_attr<ccl::reduce_scatter_attr>();
    }

    template <class Dtype, class... Args>
    void start_internal(ccl::communicator& comm,
                        size_t send_count,
                        const Dtype send_buf,
                        Dtype recv_buf,
                        const bench_exec_attr& bench_attr,
                        req_list_t& reqs,
                        Args&&... args) {
        size_t recv_count = send_count / comm.size();

        if (recv_count == 0) {
            reqs.push_back(ccl::event());
            return;
        }

        auto recv_buf_adjusted = recv_buf;
        if (recv_buf == nullptr) {
            // In-place operation
            size_t offset_count = recv_count * comm.rank();
            recv_buf_adjusted = send_buf + offset_count;
        }

        reqs.push_back(ccl::reduce_scatter(send_buf,
                                           recv_buf_adjusted,
                                           recv_count,
                                           get_ccl_dtype<Dtype>(),
                                           bench_attr.reduction,
                                           comm,
                                           std::forward<Args>(args)...));
    }
};
