//===----------------------------------------------------------------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES.
//
//===----------------------------------------------------------------------===//

#ifndef _CUDA___NVTX_NVTX_H
#define _CUDA___NVTX_NVTX_H

#include <cuda/std/detail/__config>

#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
#  pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
#  pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
#  pragma system_header
#endif // no system header

#ifdef _CCCL_DOXYGEN_INVOKED // Only parse this during doxygen passes:
//! When this macro is defined, no NVTX ranges are emitted by CCCL
#  define CCCL_DISABLE_NVTX
#endif // _CCCL_DOXYGEN_INVOKED

#define _CCCL_HAS_NVTX3() 0

// Enable the functionality of this header if:
// * The NVTX3 C API is available in CTK
// * NVTX is not explicitly disabled (via CCCL_DISABLE_NVTX or NVTX_DISABLE)
// * the compiler is not nvc++ (NVTX3 uses module as an identifier, which trips up NVHPC, fixed in CTK >= 13.0)
// * the compiler is not NVRTC
#if __has_include(<nvtx3/nvToolsExt.h>) && !defined(CCCL_DISABLE_NVTX) && !defined(NVTX_DISABLE) \
                      && (!_CCCL_COMPILER(NVHPC) || _CCCL_CTK_AT_LEAST(13, 0))                       \
                      && !_CCCL_COMPILER(NVRTC)

// Since NVTX 3.2, the NVTX headers can declare themselves as system headers by declaring the following macro:
#  ifdef NVTX_AS_SYSTEM_HEADER
#    define NVTX_AS_SYSTEM_HEADER_DEFINED_BY_USER
#  else // NVTX_AS_SYSTEM_HEADER
#    define NVTX_AS_SYSTEM_HEADER
#  endif // NVTX_AS_SYSTEM_HEADER

// Include our NVTX3 C++ wrapper if not available from the CTK or not provided by the user
// Note: NVTX3 is available in the CTK since 12.9, so we can drop our copy once this is the minimum supported version
#  if __has_include(<nvtx3/nvtx3.hpp>)
#    include <nvtx3/nvtx3.hpp>
#  else // __has_include(<nvtx3/nvtx3.hpp>)
#    include <cuda/__nvtx/nvtx3.h>
#  endif // __has_include(<nvtx3/nvtx3.hpp>)

#  ifndef NVTX_AS_SYSTEM_HEADER_DEFINED_BY_USER
#    undef NVTX_AS_SYSTEM_HEADER
#  endif // NVTX_AS_SYSTEM_HEADER_DEFINED_BY_USER
#  undef NVTX_AS_SYSTEM_HEADER_DEFINED_BY_USER

// We expect the NVTX3 V1 C++ API to be available when nvtx3.hpp is available. This should work, because newer versions
// of NVTX3 will continue to declare previous API versions. See also:
// https://github.com/NVIDIA/NVTX/blob/release-v3/c/include/nvtx3/nvtx3.hpp#L2835-L2841.
#  ifdef NVTX3_CPP_DEFINITIONS_V1_0
#    undef _CCCL_HAS_NVTX3
#    define _CCCL_HAS_NVTX3() 1
#  else // NVTX3_CPP_DEFINITIONS_V1_0
// If this happens NVTX3 changed in a way we did not anticipate, and we need to get in touch with them
#    if _CCCL_COMPILER(MSVC)
#      pragma message( \
        "warning: nvtx3.h is available but does not define the V1 API. This is odd. Please open a GitHub issue at: https://github.com/NVIDIA/cccl/issues.")
#    else
#      warning nvtx3.h is available but does not define the V1 API. This is odd. Please open a GitHub issue at: https://github.com/NVIDIA/cccl/issues.
#    endif
#  endif // NVTX3_CPP_DEFINITIONS_V1_0
#endif // __has_include(<nvtx3/nvToolsExt.h>) && !defined(CCCL_DISABLE_NVTX) && !defined(NVTX_DISABLE) &&
       // (!_CCCL_COMPILER(NVHPC)) && !_CCCL_COMPILER(NVRTC)

#if _CCCL_HAS_NVTX3()
#  include <cuda/std/__cccl/prologue.h>

_CCCL_BEGIN_NAMESPACE_CUDA
struct __nvtx_cccl_domain
{
  static constexpr const char* name{"CCCL"};
};

using __nvtx_cccl_range = ::nvtx3::v1::scoped_range_in<__nvtx_cccl_domain>;

// this type ensures that no NVTX range code is emitted in device code
struct __nvtx_cccl_optional_range_host_only
{
  bool __engaged = false;
  alignas(__nvtx_cccl_range) unsigned char __storage[sizeof(__nvtx_cccl_range)];

  __nvtx_cccl_optional_range_host_only() = default;

  _CCCL_HOST_API void __start(const ::nvtx3::v1::event_attributes& __attributes)
  {
    ::new (__storage) __nvtx_cccl_range(__attributes);
    __engaged = true;
  }

  _CCCL_API ~__nvtx_cccl_optional_range_host_only()
  {
    NV_IF_TARGET(NV_IS_HOST, ({
                   if (__engaged)
                   {
                     reinterpret_cast<__nvtx_cccl_range*>(__storage)->~__nvtx_cccl_range();
                   }
                 }));
  }
};
_CCCL_END_NAMESPACE_CUDA

// Hook for the NestedNVTXRangeGuard from the unit tests
#  ifndef _CCCL_BEFORE_NVTX_RANGE_SCOPE
#    define _CCCL_BEFORE_NVTX_RANGE_SCOPE(name)
#  endif // !CCCL_DETAIL_BEFORE_NVTX_RANGE_SCOPE

#  if _CCCL_HOST_COMPILATION()
// Conditionally inserts a NVTX range starting here until the end of the current function scope in host code. Does
// nothing in device code.
// The __nvtx_cccl_optional_range_host_only type (a simplified optional<T>) is needed to defer the construction of the
// NVTX range and message string registration (static variables) into a region running only on the host, while
// preserving the semantic scope where the range is declared.
#    define _CCCL_NVTX_RANGE_SCOPE_IF(condition, name)                                                             \
      _CCCL_BEFORE_NVTX_RANGE_SCOPE(name)                                                                          \
      ::cuda::__nvtx_cccl_optional_range_host_only __cuda_nvtx3_range;                                             \
      NV_IF_TARGET(                                                                                                \
        NV_IS_HOST, ({                                                                                             \
          static const ::nvtx3::v1::registered_string_in<::cuda::__nvtx_cccl_domain> __cuda_nvtx3_func_name{name}; \
          static const ::nvtx3::v1::event_attributes __cuda_nvtx3_func_attr{__cuda_nvtx3_func_name};               \
          if (condition)                                                                                           \
          {                                                                                                        \
            __cuda_nvtx3_range.__start(__cuda_nvtx3_func_attr);                                                    \
          }                                                                                                        \
        }))
#  else // ^^^ _CCCL_HOST_COMPILATION() ^^^ / vvv !_CCCL_HOST_COMPILATION() vvv
#    define _CCCL_NVTX_RANGE_SCOPE_IF(condition, name)
#  endif // ^^^ !_CCCL_HOST_COMPILATION() ^^^

#  define _CCCL_NVTX_RANGE_SCOPE(name) _CCCL_NVTX_RANGE_SCOPE_IF(true, name)

#  include <cuda/std/__cccl/epilogue.h>

#else // _CCCL_HAS_NVTX3()
#  define _CCCL_NVTX_RANGE_SCOPE_IF(condition, name)
#  define _CCCL_NVTX_RANGE_SCOPE(name)
#endif // _CCCL_HAS_NVTX3()

#endif // _CUDA___NVTX_NVTX_H
