/*------------------------------------------------------------------------------------------------*
 * Copyright (C) by the DBCSR developers group - All rights reserved                              *
 * This file is part of the DBCSR library.                                                        *
 *                                                                                                *
 * For information on the license, see the LICENSE file.                                          *
 * For further information please visit https://dbcsr.cp2k.org                                    *
 * SPDX-License-Identifier: GPL-2.0+                                                              *
 *------------------------------------------------------------------------------------------------*/

#if defined( __CUDA_PROFILING)

#include <nvToolsExt.h>
#include <stdio.h>
#include <pthread.h>

const uint32_t colormap[] = { 0xFFFFFF00,  // Yellow
                              0xFFFF00FF,  // Fuchsia
                              0xFFFF0000,  // Red
                              0xFFC0C0C0,  // Silver
                              0xFF808080,  // Gray
                              0xFF808000,  // Olive
                              0xFF800080,  // Purple
                              0xFF800000,  // Maroon
                              0xFF00FFFF,  // Aqua
                              0xFF00FF00,  // Lime
                              0xFF008080,  // Teal
                              0xFF008000,  // Green
                              0xFF0000FF,  // Blue
                              0xFF000080}; // Navy

//==============================================================================
extern "C" int cuda_nvtx_range_push_cu(const char* message) {

    // assembling event attribute
    nvtxEventAttributes_t eventAttrib = {0};
    eventAttrib.version = NVTX_VERSION;
    eventAttrib.size = NVTX_EVENT_ATTRIB_STRUCT_SIZE;
    eventAttrib.messageType = NVTX_MESSAGE_TYPE_ASCII;
    eventAttrib.message.ascii = message;

    // colors are picked based on a (very simple) hash value of the message
    int hash=0;
    for (int i=0; i < strlen(message); i++)
        hash += i*message[i]*message[i];
    eventAttrib.colorType = NVTX_COLOR_ARGB;
    eventAttrib.color = colormap[hash%14];

    // these field could be filled with useful stuff
    eventAttrib.payloadType = NVTX_PAYLOAD_TYPE_INT64;
    eventAttrib.payload.llValue = 123;
    eventAttrib.category = 42;

    int level = nvtxRangePushEx(&eventAttrib);
    return(level);
}

//==============================================================================
extern "C" int cuda_nvtx_range_pop_cu() {
    int level = nvtxRangePop();
    return(level);
}

//==============================================================================
extern "C" void cuda_nvtx_name_osthread_cu(char* name){
    nvtxNameOsThread(pthread_self(), name);
}

#endif
