BertanDogancay
178 строки
7.2 KiB
C++
178 строки
7.2 KiB
C++
/*************************************************************************
|
|
* Copyright (c) 2019-2020, NVIDIA CORPORATION. All rights reserved.
|
|
*
|
|
* See LICENSE.txt for license information
|
|
************************************************************************/
|
|
|
|
#ifndef NCCL_NVTX_H_
|
|
#define NCCL_NVTX_H_
|
|
|
|
#include "nvtx3/nvtx3.hpp"
|
|
#include "roctx.h"
|
|
|
|
#if __cpp_constexpr >= 201304L && !defined(NVTX3_CONSTEXPR_IF_CPP14)
|
|
#define NVTX3_CONSTEXPR_IF_CPP14 constexpr
|
|
#else
|
|
#define NVTX3_CONSTEXPR_IF_CPP14
|
|
#endif
|
|
|
|
// Define all NCCL-provided static schema IDs here (avoid duplicates).
|
|
#define NVTX_SID_CommInitRank 0
|
|
#define NVTX_SID_CommInitAll 1
|
|
#define NVTX_SID_CommDestroy 2 // same schema as NVTX_SID_CommInitRank
|
|
#define NVTX_SID_CommAbort 3 // same schema as NVTX_SID_CommInitRank
|
|
#define NVTX_SID_AllGather 4
|
|
#define NVTX_SID_AllReduce 5
|
|
#define NVTX_SID_AllToAll 6
|
|
#define NVTX_SID_AllToAllv 7
|
|
#define NVTX_SID_Broadcast 8
|
|
#define NVTX_SID_Gather 9
|
|
#define NVTX_SID_MSCCL 10
|
|
#define NVTX_SID_ReduceScatter 11
|
|
#define NVTX_SID_Reduce 12
|
|
#define NVTX_SID_Scatter 13
|
|
#define NVTX_SID_Send 14
|
|
#define NVTX_SID_Recv 15
|
|
#define NVTX_SID_CommInitRankConfig 16 // same schema as NVTX_SID_CommInitRank
|
|
#define NVTX_SID_CommInitRankScalable 17 // same schema as NVTX_SID_CommInitRank
|
|
#define NVTX_SID_CommSplit 18
|
|
#define NVTX_SID_CommFinalize 19
|
|
// When adding new schema IDs, DO NOT re-use/overlap with the enum schema ID below!
|
|
|
|
// Define static schema ID for the reduction operation.
|
|
#define NVTX_PAYLOAD_ENTRY_NCCL_REDOP 20 + NVTX_PAYLOAD_ENTRY_TYPE_SCHEMA_ID_STATIC_START
|
|
|
|
extern const nvtxDomainHandle_t ncclNvtxDomainHandle;
|
|
|
|
struct nccl_domain{static constexpr char const* name{"NCCL"};};
|
|
|
|
/// @brief Register an NVTX payload schema for static-size payloads.
|
|
class payload_schema {
|
|
public:
|
|
explicit payload_schema(const nvtxPayloadSchemaEntry_t entries[], size_t numEntries,
|
|
const uint64_t schemaId, const size_t size) noexcept
|
|
{
|
|
schema_attr.payloadStaticSize = size;
|
|
schema_attr.entries = entries;
|
|
schema_attr.numEntries = numEntries;
|
|
schema_attr.schemaId = schemaId;
|
|
nvtxPayloadSchemaRegister(nvtx3::domain::get<nccl_domain>(), &schema_attr);
|
|
}
|
|
|
|
payload_schema() = delete;
|
|
~payload_schema() = default;
|
|
payload_schema(payload_schema const&) = default;
|
|
payload_schema& operator=(payload_schema const&) = default;
|
|
payload_schema(payload_schema&&) = default;
|
|
payload_schema& operator=(payload_schema&&) = default;
|
|
|
|
private:
|
|
nvtxPayloadSchemaAttr_t schema_attr{
|
|
NVTX_PAYLOAD_SCHEMA_ATTR_TYPE |
|
|
NVTX_PAYLOAD_SCHEMA_ATTR_ENTRIES |
|
|
NVTX_PAYLOAD_SCHEMA_ATTR_NUM_ENTRIES |
|
|
NVTX_PAYLOAD_SCHEMA_ATTR_STATIC_SIZE |
|
|
NVTX_PAYLOAD_SCHEMA_ATTR_SCHEMA_ID,
|
|
nullptr, /* schema name is not needed */
|
|
NVTX_PAYLOAD_SCHEMA_TYPE_STATIC,
|
|
NVTX_PAYLOAD_SCHEMA_FLAG_NONE,
|
|
nullptr, 0, 0, 0, 0, nullptr};
|
|
};
|
|
|
|
// Convenience macro to give the payload parameters a scope.
|
|
#define NVTX3_PAYLOAD(...) __VA_ARGS__
|
|
|
|
// Create NVTX push/pop range with parameters
|
|
// @param N NCCL API name without the `nccl` prefix.
|
|
// @param T name of the used NVTX payload schema without "Schema" suffix.
|
|
// @param P payload parameters/entries
|
|
#if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__)
|
|
#define NVTX3_FUNC_WITH_PARAMS(N, T, P) \
|
|
constexpr uint64_t schemaId = NVTX_PAYLOAD_ENTRY_TYPE_SCHEMA_ID_STATIC_START + NVTX_SID_##N; \
|
|
static const payload_schema schema{T##Schema, std::extent<decltype(T##Schema)>::value - 1, \
|
|
schemaId, sizeof(T)}; \
|
|
const T _payload = {P}; \
|
|
nvtxPayloadData_t nvtx3_bpl__[] = {{schemaId, sizeof(_payload), &_payload}}; \
|
|
roctx_scoped_range_in const roctx_range__{T##Schema, nvtx3_bpl__, std::extent<decltype(T##Schema)>::value - 1, "RCCL_" #N};
|
|
#else
|
|
#define NVTX3_FUNC_WITH_PARAMS(N, T, P) \
|
|
constexpr uint64_t schemaId = NVTX_PAYLOAD_ENTRY_TYPE_SCHEMA_ID_STATIC_START + NVTX_SID_##N; \
|
|
static const payload_schema schema{T##Schema, std::extent<decltype(T##Schema)>::value - 1, \
|
|
schemaId, sizeof(T)}; \
|
|
static ::nvtx3::v1::registered_string_in<nccl_domain> const nvtx3_func_name__{__func__}; \
|
|
const T _payload = {P}; \
|
|
nvtxPayloadData_t nvtx3_bpl__[] = {{schemaId, sizeof(_payload), &_payload}}; \
|
|
::nvtx3::v1::event_attributes const nvtx3_func_attr__{nvtx3_func_name__, nvtx3_bpl__}; \
|
|
::nvtx3::v1::scoped_range_in<nccl_domain> const nvtx3_range__{nvtx3_func_attr__};
|
|
#endif
|
|
|
|
/// @brief Creates an NVTX range with extended payload using the RAII pattern.
|
|
/// @tparam PayloadType Data type of the payload.
|
|
template <typename PayloadType>
|
|
class ncclNvtxRange {
|
|
public:
|
|
explicit ncclNvtxRange(const nvtxEventAttributes_t* evtAttr) noexcept {
|
|
nvtxDomainRangePushEx(nvtx3::domain::get<nccl_domain>(), evtAttr);
|
|
}
|
|
|
|
~ncclNvtxRange() noexcept {
|
|
if (payloadData.payload) {
|
|
nvtxRangePopPayload(nvtx3::domain::get<nccl_domain>(), &payloadData, 1);
|
|
} else {
|
|
nvtxDomainRangePop(nvtx3::domain::get<nccl_domain>());
|
|
}
|
|
}
|
|
|
|
void setPayloadData(const uint64_t schemaId) noexcept
|
|
{
|
|
payloadData = {schemaId, sizeof(PayloadType), &payload};
|
|
}
|
|
|
|
ncclNvtxRange() = delete;
|
|
ncclNvtxRange(ncclNvtxRange const&) = default;
|
|
ncclNvtxRange& operator=(ncclNvtxRange const&) = default;
|
|
ncclNvtxRange(ncclNvtxRange&&) = default;
|
|
ncclNvtxRange& operator=(ncclNvtxRange&&) = default;
|
|
|
|
// Holds the payload data.
|
|
PayloadType payload{};
|
|
|
|
nvtxPayloadData_t payloadData = {NVTX_PAYLOAD_ENTRY_TYPE_INVALID, 0, NULL};
|
|
};
|
|
|
|
// Create an NVTX range with the function name as the range name. Use RAII pattern.
|
|
// @param T Type ID of the NVTX payload (pointer for variable-size payloads).
|
|
#define NVTX3_RANGE(T) \
|
|
static ::nvtx3::v1::registered_string_in<nccl_domain> const nvtx3_func_name__{__func__}; \
|
|
::nvtx3::v1::event_attributes const nvtx3_func_attr__{nvtx3_func_name__}; \
|
|
ncclNvtxRange<T> nvtx3_range__{nvtx3_func_attr__.get()};
|
|
|
|
// Add static-size payload to the NVTX range created with `NVTX3_RANGE()`,
|
|
// which must be in this or an outer scope.
|
|
// @param N NCCL API name without the `nccl` prefix.
|
|
// @param S name of the used NVTX payload schema.
|
|
// @param P payload parameters/entries
|
|
#if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__)
|
|
#define NVTX3_RANGE_ADD_PAYLOAD(N, S, P) do { \
|
|
constexpr uint64_t schema_id = NVTX_PAYLOAD_ENTRY_TYPE_SCHEMA_ID_STATIC_START + NVTX_SID_##N; \
|
|
static const payload_schema schema{S, std::extent<decltype(S)>::value - 1, schema_id, \
|
|
sizeof(nvtx3_range__.payload)}; \
|
|
nvtx3_range__.payload = {P}; \
|
|
nvtx3_range__.setPayloadData(schema_id); \
|
|
nvtxPayloadData_t nvtx3_bpl__[] = {{schema_id, sizeof(nvtx3_range__.payloadData), &nvtx3_range__.payloadData}}; \
|
|
roctx_scoped_range_in const roctx_range__{S, nvtx3_bpl__, std::extent<decltype(S)>::value - 1, "RCCL_" #N}; \
|
|
} while (0)
|
|
#else
|
|
#define NVTX3_RANGE_ADD_PAYLOAD(N, S, P) do { \
|
|
constexpr uint64_t schema_id = NVTX_PAYLOAD_ENTRY_TYPE_SCHEMA_ID_STATIC_START + NVTX_SID_##N; \
|
|
static const payload_schema schema{S, std::extent<decltype(S)>::value - 1, schema_id, \
|
|
sizeof(nvtx3_range__.payload)}; \
|
|
nvtx3_range__.payload = {P}; \
|
|
nvtx3_range__.setPayloadData(schema_id); \
|
|
} while (0)
|
|
#endif
|
|
|
|
extern void initNvtxRegisteredEnums();
|
|
|
|
#endif
|