// SPDX-License-Identifier: BSD-2-Clause
/* Copyright (C) 2017 - 2021 Intel Corporation. */

#include "allocator_perf_tool/Allocation_info.hpp"
#include "allocator_perf_tool/GTestAdapter.hpp"
#include <gtest/gtest.h>
#include <hbwmalloc.h>
#include <memkind.h>
#include <memory>
#include <numa.h>
#include <stdio.h>
#include <vector>

typedef std::unique_ptr<void, void (*)(void *)> hbw_mem_ptr;

class HBWPreferredLocalityTest: public ::testing::Test
{
private:
    int find_closest_node(int node, const std::vector<int> &nodes)
    {
        int min_distance = 0;
        int closest_node = -1;
        for (int i = 0; i < nodes.size(); i++) {
            int distance = numa_distance(node, nodes[i]);
            if (distance && (distance < min_distance || min_distance == 0)) {
                min_distance = distance;
                closest_node = nodes[i];
            }
        }
        return closest_node;
    }

    bool pin_to_cpu(int cpu_id)
    {
        cpu_set_t cpu_set;
        CPU_ZERO(&cpu_set);
        CPU_SET(cpu_id, &cpu_set);
        return sched_setaffinity(0, sizeof(cpu_set_t), &cpu_set) != -1;
    }

    void check_ptr_numa(void *ptr, int expected_numa_id, int cpu_id,
                        size_t size)
    {
        memset(ptr, 1, size);

        int numa_id = get_numa_node_id(ptr);
        EXPECT_EQ(numa_id, expected_numa_id);

        char property_name[50];
        snprintf(property_name, 50, "actual_numa_for_cpu_%d_expected_numa_%d",
                 cpu_id, expected_numa_id);
        GTestAdapter::RecordProperty(property_name, numa_id);
    }

public:
    void
    pin_memory_in_requesting_mem_thread(size_t size,
                                        const std::vector<int> &cpu_ids,
                                        const std::vector<int> &mcdram_nodes)
    {
#ifdef _OPENMP
        int threads_num = cpu_ids.size();
        int ret = hbw_set_policy(HBW_POLICY_PREFERRED);
        ASSERT_EQ(ret, 0);

#pragma omp parallel for num_threads(threads_num)
        for (int i = 0; i < threads_num; i++) {
            if (!pin_to_cpu(cpu_ids[i])) {
                ADD_FAILURE();
                continue;
            }

            void *internal_ptr = hbw_malloc(size);
            if (!internal_ptr) {
                ADD_FAILURE();
                continue;
            }
            hbw_mem_ptr ptr(internal_ptr, hbw_free);
            int expected_numa_id =
                find_closest_node(numa_node_of_cpu(cpu_ids[i]), mcdram_nodes);
            check_ptr_numa(ptr.get(), expected_numa_id, cpu_ids[i], size);
        }
#else
        std::cout << "[ SKIPPED ] Feature OPENMP not supported" << std::endl;
#endif
    }

    void pin_memory_in_other_thread_than_requesting_mem(
        size_t size, const std::vector<int> &cpu_ids,
        const std::vector<int> &mcdram_nodes)
    {
#ifdef _OPENMP
        int threads_num = cpu_ids.size();
        int ret = hbw_set_policy(HBW_POLICY_PREFERRED);
        ASSERT_EQ(ret, 0);

        int main_thread_cpu_id = 0;
        int expected_numa_id =
            find_closest_node(main_thread_cpu_id, mcdram_nodes);
        ASSERT_TRUE(pin_to_cpu(main_thread_cpu_id));

        std::vector<hbw_mem_ptr> ptrs;
        for (int i = 0; i < threads_num; i++) {
            void *internal_ptr = hbw_malloc(size);
            ASSERT_TRUE(internal_ptr);
            ptrs.emplace_back(internal_ptr, hbw_free);
        }

#pragma omp parallel for num_threads(threads_num)
        for (int i = 0; i < threads_num; i++) {
            if (!pin_to_cpu(cpu_ids[i])) {
                ADD_FAILURE();
                continue;
            }
            check_ptr_numa(ptrs[i].get(), expected_numa_id, cpu_ids[i], size);
        }
#else
        std::cout << "[ SKIPPED ] Feature OPENMP not supported" << std::endl;
#endif
    }
};

TEST_F(
    HBWPreferredLocalityTest,
    test_TC_MEMKIND_KNL_SNC4_pin_memory_in_requesting_mem_thread_4_threads_100_bytes)
{
    pin_memory_in_requesting_mem_thread(100u, std::vector<int>{0, 18, 36, 54},
                                        std::vector<int>{4, 5, 6, 7});
}

TEST_F(
    HBWPreferredLocalityTest,
    test_TC_MEMKIND_KNL_SNC4_pin_memory_in_other_thread_than_requesting_mem_4_threads_100_bytes)
{
    pin_memory_in_other_thread_than_requesting_mem(
        100u, std::vector<int>{0, 18, 36, 54}, std::vector<int>{4, 5, 6, 7});
}
