gpu/rocm.hpp

Go to the documentation of this file.

// Copyright (c) 2014-2023, Lawrence Livermore National Security, LLC.
// Produced at the Lawrence Livermore National Laboratory.
// Written by the LBANN Research Team (B. Van Essen, et al.) listed in
// the CONTRIBUTORS file. <lbann-dev@llnl.gov>
//
// LLNL-CODE-697807.
// All rights reserved.
//
// This file is part of LBANN: Livermore Big Artificial Neural Network
// Toolkit. For details, see http://software.llnl.gov/LBANN or
// https://github.com/LLNL/LBANN.
//
// Licensed under the Apache License, Version 2.0 (the "Licensee"); 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.

#ifndef LBANN_UTILS_ROCM_HPP
#define LBANN_UTILS_ROCM_HPP

#include "lbann/utils/exception.hpp"
#include "lbann_config.hpp"

#ifdef LBANN_HAS_ROCM

#include <hip/hip_runtime.h>
#include <thrust/detail/allocator/tagged_allocator.h>
#include <thrust/device_vector.h>
#include <thrust/memory.h>
#include <thrust/system/hip/detail/par.h>
#include <thrust/version.h>

// -------------------------------------------------------------
// Error utility macros
// -------------------------------------------------------------
#define LBANN_ROCM_SYNC(async)                                                 \
  do {                                                                         \
    /* Synchronize GPU and check for errors. */                                \
    hipError_t status_ROCM_SYNC = hipDeviceSynchronize();                      \
    if (status_ROCM_SYNC == hipSuccess)                                        \
      status_ROCM_SYNC = hipGetLastError();                                    \
    if (status_ROCM_SYNC != hipSuccess) {                                      \
      LBANN_ERROR((async ? "Asynchronous " : ""),                              \
                  "ROCm error (",                                              \
                  hipGetErrorString(status_ROCM_SYNC),                         \
                  ")");                                                        \
    }                                                                          \
  } while (0)
#define LBANN_ROCM_CHECK_LAST_ERROR(async)                                     \
  do {                                                                         \
    hipError_t status = hipGetLastError();                                     \
    if (status != hipSuccess) {                                                \
      LBANN_ERROR((async ? "Asynchronous " : ""),                              \
                  "ROCm error (",                                              \
                  hipGetErrorString(status),                                   \
                  ")");                                                        \
    }                                                                          \
  } while (0)
#define FORCE_CHECK_ROCM(rocm_call)                                            \
  do {                                                                         \
    /* Call ROCM API routine, synchronizing before and */                      \
    /* after to check for errors. */                                           \
    LBANN_ROCM_SYNC(true);                                                     \
    hipError_t status_CHECK_ROCM = (rocm_call);                                \
    if (status_CHECK_ROCM != hipSuccess) {                                     \
      LBANN_ERROR("ROCm error (", hipGetErrorString(status_CHECK_ROCM), ")");  \
    }                                                                          \
    LBANN_ROCM_SYNC(false);                                                    \
  } while (0)
#define FORCE_CHECK_ROCM_NOSYNC(rocm_call)                                     \
  do {                                                                         \
    hipError_t status_CHECK_ROCM = (rocm_call);                                \
    if (status_CHECK_ROCM != hipSuccess) {                                     \
      LBANN_ERROR("ROCm error (", hipGetErrorString(status_CHECK_ROCM), ")");  \
    }                                                                          \
  } while (0)
#ifdef LBANN_DEBUG
#define CHECK_ROCM(rocm_call) FORCE_CHECK_ROCM(rocm_call);
#else
#define CHECK_ROCM(rocm_call) FORCE_CHECK_ROCM_NOSYNC(rocm_call)
#endif // #ifdef LBANN_DEBUG

namespace lbann {
namespace rocm {

constexpr hipMemcpyKind GPU_MEMCPY_DEVICE_TO_DEVICE = hipMemcpyDeviceToDevice;

// -------------------------------------------------------------
// Wrapper classes
// -------------------------------------------------------------

class event_wrapper
{
public:
  event_wrapper();
  event_wrapper(const event_wrapper& other);
  event_wrapper& operator=(const event_wrapper& other);
  ~event_wrapper();
  void record(hipStream_t stream);
  bool query() const;
  void synchronize();
  hipEvent_t& get_event();

private:
  hipEvent_t m_event;
  hipStream_t m_stream;
};

// -------------------------------------------------------------
// Helper functions for tensor operations
// -------------------------------------------------------------

template <typename TensorDataType>
void copy_tensor(hipStream_t stream,
                 const std::vector<size_t>& dims,
                 const TensorDataType* input,
                 const std::vector<size_t>& input_strides,
                 TensorDataType* output,
                 const std::vector<size_t>& output_strides);

void mem_copy_async(void* output,
                    const void* input,
                    const size_t count,
                    hipMemcpyKind kind,
                    hipStream_t stream);

// -------------------------------------------------------------
// Utilities for Thrust
// -------------------------------------------------------------
namespace thrust {

using execute_on_stream = ::thrust::hip_rocprim::execute_on_stream;

template <typename T = El::byte>
class allocator : public ::thrust::detail::tagged_allocator<
                    T,
                    execute_on_stream,
                    ::thrust::pointer<T, execute_on_stream>>
{
public:
  // Convenient typedefs
  typedef ::thrust::detail::tagged_allocator<
    T,
    execute_on_stream,
    ::thrust::pointer<T, execute_on_stream>>
    parent_class;
  typedef typename parent_class::value_type value_type;
  typedef typename parent_class::pointer pointer;
  typedef typename parent_class::size_type size_type;
  typedef typename parent_class::system_type system_type;

  allocator(hipStream_t stream = hydrogen::rocm::GetDefaultStream());
  pointer allocate(size_type size);
  void deallocate(pointer buffer, size_type size = 0);
  system_type& system();

private:
  hipStream_t m_stream;
  system_type m_system;
};

template <typename T>
using vector = ::thrust::device_vector<T, allocator<T>>;

} // namespace thrust
} // namespace rocm
} // namespace lbann

#endif // LBANN_HAS_ROCM
#endif // LBANN_UTILS_ROCM_HPP