Line data    Source code

       1             : /*----------------------------------------------------------------------------*/
       2             : /*  CP2K: A general program to perform molecular dynamics simulations         */
       3             : /*  Copyright 2000-2022 CP2K developers group <https://cp2k.org>              */
       4             : /*                                                                            */
       5             : /*  SPDX-License-Identifier: BSD-3-Clause                                     */
       6             : /*----------------------------------------------------------------------------*/
       7             : 
       8             : #include <assert.h>
       9             : #include <stdint.h>
      10             : #include <stdio.h>
      11             : #include <stdlib.h>
      12             : #include <string.h>
      13             : 
      14             : #include "offload_library.h"
      15             : #include "offload_runtime.h"
      16             : 
      17             : #if defined(__OFFLOAD_PROFILING)
      18             : #if defined(__OFFLOAD_CUDA)
      19             : #include <nvToolsExt.h>
      20             : #elif defined(__OFFLOAD_HIP) && defined(__HIP_PLATFORM_AMD__)
      21             : #include <roctracer/roctx.h>
      22             : #endif
      23             : #endif
      24             : 
      25             : static int chosen_device_id = -1;
      26             : 
      27             : const uint32_t colormap[] = {0xFFFFFF00,  // Yellow
      28             :                              0xFFFF00FF,  // Fuchsia
      29             :                              0xFFFF0000,  // Red
      30             :                              0xFFC0C0C0,  // Silver
      31             :                              0xFF808080,  // Gray
      32             :                              0xFF808000,  // Olive
      33             :                              0xFF800080,  // Purple
      34             :                              0xFF800000,  // Maroon
      35             :                              0xFF00FFFF,  // Aqua
      36             :                              0xFF00FF00,  // Lime
      37             :                              0xFF008080,  // Teal
      38             :                              0xFF008000,  // Green
      39             :                              0xFF0000FF,  // Blue
      40             :                              0xFF000080}; // Navy
      41             : 
      42             : /*******************************************************************************
      43             :  * \brief Returns the number of available devices.
      44             :  * \author Ole Schuett
      45             :  ******************************************************************************/
      46        7610 : int offload_get_device_count(void) {
      47        7610 :   int count = 0;
      48             : #ifdef __OFFLOAD_CUDA
      49             :   OFFLOAD_CHECK(cudaGetDeviceCount(&count));
      50             : #elif defined(__OFFLOAD_HIP)
      51             :   OFFLOAD_CHECK(hipGetDeviceCount(&count));
      52             : #endif
      53        7610 :   return count;
      54             : }
      55             : 
      56             : /*******************************************************************************
      57             :  * \brief Selects the chosen device to be used.
      58             :  * \author Ole Schuett
      59             :  ******************************************************************************/
      60           2 : void offload_set_chosen_device(int device_id) { chosen_device_id = device_id; }
      61             : 
      62             : /*******************************************************************************
      63             :  * \brief Returns the chosen device.
      64             :  * \author Ole Schuett
      65             :  ******************************************************************************/
      66           0 : int offload_get_chosen_device(void) { return chosen_device_id; }
      67             : 
      68             : /*******************************************************************************
      69             :  * \brief Activates the device selected via offload_set_chosen_device()
      70             :  * \author Ole Schuett
      71             :  ******************************************************************************/
      72      993458 : void offload_activate_chosen_device(void) {
      73             : #ifdef __OFFLOAD_CUDA
      74             :   OFFLOAD_CHECK(cudaSetDevice(chosen_device_id));
      75             : #elif defined(__OFFLOAD_HIP)
      76             :   OFFLOAD_CHECK(hipSetDevice(chosen_device_id));
      77             : #endif
      78      993458 : }
      79             : 
      80             : /*******************************************************************************
      81             :  * \brief Starts a timing range.
      82             :  * \author Ole Schuett
      83             :  ******************************************************************************/
      84  1336531011 : void offload_timeset(const char *message) {
      85             : #if defined(__OFFLOAD_PROFILING)
      86             : #if defined(__OFFLOAD_CUDA)
      87             :   // colors are picked based on a (very simple) hash value of the message
      88             :   int hash = 0;
      89             :   for (size_t i = 0; i < strlen(message); i++) {
      90             :     hash += i * message[i] * message[i];
      91             :   }
      92             :   nvtxEventAttributes_t eventAttrib = {0};
      93             :   eventAttrib.version = NVTX_VERSION;
      94             :   eventAttrib.size = NVTX_EVENT_ATTRIB_STRUCT_SIZE;
      95             :   eventAttrib.messageType = NVTX_MESSAGE_TYPE_ASCII;
      96             :   eventAttrib.message.ascii = message;
      97             :   eventAttrib.colorType = NVTX_COLOR_ARGB;
      98             :   eventAttrib.color = colormap[hash % 14];
      99             :   eventAttrib.payloadType = NVTX_PAYLOAD_TYPE_INT64;
     100             :   eventAttrib.payload.llValue = 123;
     101             :   eventAttrib.category = 42;
     102             :   nvtxRangePushEx(&eventAttrib);
     103             : #elif defined(__OFFLOAD_HIP) && defined(__HIP_PLATFORM_AMD__)
     104             :   roctxRangePushA(message);
     105             : #endif
     106             : #endif
     107  1336531011 :   (void)message; // mark argument as used
     108  1336531011 : }
     109             : 
     110             : /*******************************************************************************
     111             :  * \brief Ends a timing range.
     112             :  * \author Ole Schuett
     113             :  ******************************************************************************/
     114  1336531011 : void offload_timestop(void) {
     115             : #if defined(__OFFLOAD_PROFILING)
     116             : #if defined(__OFFLOAD_CUDA)
     117             :   nvtxRangePop();
     118             : #elif defined(__OFFLOAD_HIP) && defined(__HIP_PLATFORM_AMD__)
     119             :   roctxRangePop();
     120             : #endif
     121             : #endif
     122  1336531011 : }
     123             : 
     124             : /*******************************************************************************
     125             :  * \brief Gets free and total device memory.
     126             :  * \author Ole Schuett
     127             :  ******************************************************************************/
     128      197833 : void offload_mem_info(size_t *free, size_t *total) {
     129             : #if defined(__OFFLOAD_CUDA)
     130             :   OFFLOAD_CHECK(cudaMemGetInfo(free, total));
     131             : #elif defined(__OFFLOAD_HIP)
     132             :   OFFLOAD_CHECK(hipMemGetInfo(free, total));
     133             : #else
     134      197833 :   *free = 0;
     135      197833 :   *total = 0;
     136             : #endif
     137      197833 : }
     138             : 
     139           0 : int offload_host_malloc(void **ptr__, const size_t size__) {
     140             : #if defined(__OFFLOAD_CUDA) || defined(__OFFLOAD_HIP)
     141             :   /* API checks are included in the overloading of the function */
     142             :   offloadMallocHost(ptr__, size__);
     143             : #else
     144           0 :   *ptr__ = malloc(size__);
     145           0 :   return 0;
     146             : #endif
     147             :   return 0;
     148             : }
     149             : 
     150           0 : int offload_host_free(void *ptr__) {
     151             : #if defined(__OFFLOAD_CUDA) || defined(__OFFLOAD_HIP)
     152             :   offloadFreeHost(ptr__);
     153             : #else
     154           0 :   free(ptr__);
     155             : #endif
     156           0 :   return 0;
     157             : }
     158             : 
     159             : // EOF