AI/rocm-systems
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Implement ROCTX (#1094)

Browse Source 
* Implement roctx
This commit is contained in:
Bertan Dogancay
2024-02-27 15:46:15 -07:00
committed by GitHub
parent dae6df6d16
commit b617aecc31
9 changed files with 364 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/collectives.cc
+41
View File
@@ -68,6 +68,13 @@ NCCL_API(ncclResult_t, ncclAllToAll, const void* sendbuff, void* recvbuff, size_
ncclComm_t comm, hipStream_t stream);
ncclResult_t ncclAllToAll(const void* sendbuff, void* recvbuff, size_t count, ncclDataType_t datatype,
ncclComm_t comm, hipStream_t stream) {
// Just pass the size of one message and not the total bytes sent/received.
constexpr nvtxPayloadSchemaEntry_t AllToAllSchema[] = {
{0, NVTX_PAYLOAD_ENTRY_TYPE_SIZE, "Message size [bytes]"}
};
size_t msgsize = count * ncclTypeSize(datatype);
NVTX3_FUNC_WITH_PARAMS(AllToAll, AllToAllSchema, msgsize)
if (mscclAvailable() && !mscclIsCaller()) {
return mscclEnqueueCheck(
sendbuff, nullptr, nullptr, recvbuff, nullptr, nullptr,
@@ -103,6 +110,18 @@ NCCL_API(ncclResult_t, ncclAllToAllv, const void *sendbuff, const size_t sendcou
ncclResult_t ncclAllToAllv(const void *sendbuff, const size_t sendcounts[], const size_t sdispls[],
void *recvbuff, const size_t recvcounts[], const size_t rdispls[],
ncclDataType_t datatype, ncclComm_t comm, hipStream_t stream) {
struct NvtxParamsAllToAllv {
size_t sendbytes;
size_t recvbytes;
};
// Just pass the size of one send/recv messages and not the total bytes sent/received.
constexpr nvtxPayloadSchemaEntry_t AllToAllvSchema[] = {
{0, NVTX_PAYLOAD_ENTRY_TYPE_SIZE, "Message size [bytes] (Send)"},
{0, NVTX_PAYLOAD_ENTRY_TYPE_SIZE, "Message size [bytes] (Recv)"}
};
NvtxParamsAllToAllv payload{sendcounts[comm->rank] * ncclTypeSize(datatype), recvcounts[comm->rank] * ncclTypeSize(datatype)};
NVTX3_FUNC_WITH_PARAMS(AllToAllv, AllToAllvSchema, payload)
if (mscclAvailable() && !mscclIsCaller()) {
return mscclEnqueueCheck(
sendbuff, sendcounts, sdispls, recvbuff, recvcounts, rdispls,
@@ -170,6 +189,17 @@ NCCL_API(ncclResult_t, ncclGather, const void* sendbuff, void* recvbuff, size_t
ncclDataType_t datatype, int root, ncclComm_t comm, hipStream_t stream);
ncclResult_t ncclGather(const void* sendbuff, void* recvbuff, size_t sendcount,
ncclDataType_t datatype, int root, ncclComm_t comm, hipStream_t stream) {
struct NvtxParamsGather {
size_t bytes;
int root;
};
constexpr nvtxPayloadSchemaEntry_t GatherSchema[] = {
{0, NVTX_PAYLOAD_ENTRY_TYPE_SIZE, "Bytes"},
{0, NVTX_PAYLOAD_ENTRY_TYPE_INT, "Root", nullptr, 0, offsetof(NvtxParamsGather, root)}
};
NvtxParamsGather payload{sendcount * ncclTypeSize(datatype), root};
NVTX3_FUNC_WITH_PARAMS(Gather, GatherSchema, payload)
if (mscclAvailable() && !mscclIsCaller()) {
return mscclEnqueueCheck(
sendbuff, nullptr, nullptr, recvbuff, nullptr, nullptr,
@@ -254,6 +284,17 @@ NCCL_API(ncclResult_t, ncclScatter, const void* sendbuff, void* recvbuff, size_t
ncclComm_t comm, hipStream_t stream);
ncclResult_t ncclScatter(const void* sendbuff, void* recvbuff, size_t recvcount, ncclDataType_t datatype, int root,
ncclComm_t comm, hipStream_t stream) {
struct NvtxParamsScatter {
size_t bytes;
int root;
};
constexpr nvtxPayloadSchemaEntry_t ScatterSchema[] = {
{0, NVTX_PAYLOAD_ENTRY_TYPE_SIZE, "Bytes"},
{0, NVTX_PAYLOAD_ENTRY_TYPE_INT, "Root", nullptr, 0, offsetof(NvtxParamsScatter, root)}
};
NvtxParamsScatter payload{recvcount * ncclTypeSize(datatype), root};
NVTX3_FUNC_WITH_PARAMS(Scatter, ScatterSchema, payload)
if (mscclAvailable() && !mscclIsCaller()) {
return mscclEnqueueCheck(
sendbuff, nullptr, nullptr, recvbuff, nullptr, nullptr,
src/include/nvtx.h
+18 -5
View File
@@ -8,6 +8,7 @@
#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
@@ -22,11 +23,16 @@
#define NVTX_SID_CommAbort 3 // same schema as NVTX_SID_CommInitRank
#define NVTX_SID_AllGather 4
#define NVTX_SID_AllReduce 5
#define NVTX_SID_Broadcast 6
#define NVTX_SID_ReduceScatter 7
#define NVTX_SID_Reduce 8
#define NVTX_SID_Send 9
#define NVTX_SID_Recv 10
#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 static schema ID for the reduction operation.
#define NVTX_PAYLOAD_ENTRY_NCCL_REDOP 11 + NVTX_PAYLOAD_ENTRY_TYPE_SCHEMA_ID_STATIC_START
@@ -71,6 +77,12 @@ class payload_schema {
// @param N schema name
// @param S schema (entries)
// @param P payload (struct)
#if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
#define NVTX3_FUNC_WITH_PARAMS(ID, S, P) \
nvtxPayloadData_t nvtx3_bpl__[] = { \
{NVTX_PAYLOAD_ENTRY_TYPE_SCHEMA_ID_STATIC_START + NVTX_SID_##ID, sizeof(P), &(P)}}; \
roctx_scoped_range_in const roctx_range__{S, nvtx3_bpl__, std::extent<decltype(S)>::value, "RCCL_" #ID};
#else
#define NVTX3_FUNC_WITH_PARAMS(ID, S, P) \
static const payload_schema schema{S, std::extent<decltype(S)>::value, \
NVTX_PAYLOAD_ENTRY_TYPE_SCHEMA_ID_STATIC_START + NVTX_SID_##ID, #ID}; \
@@ -79,6 +91,7 @@ class payload_schema {
{NVTX_PAYLOAD_ENTRY_TYPE_SCHEMA_ID_STATIC_START + NVTX_SID_##ID, sizeof(P), &(P)}}; \
::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
extern void initNvtxRegisteredEnums();
src/include/nvtx3/nvtx3.hpp
+5
View File
@@ -2777,10 +2777,15 @@ inline void mark(Args const&... args) noexcept
* `domain` to which the `registered_string_in` belongs. Else,
* `domain::global` to indicate that the global NVTX domain should be used.
*/
#if !defined(__HIP_PLATFORM_HCC__) && !defined(__HCC__) && !defined(__HIPCC__)
#define NVTX3_V1_FUNC_RANGE_IN(D) \
static ::nvtx3::v1::registered_string_in<D> const nvtx3_func_name__{__func__}; \
static ::nvtx3::v1::event_attributes const nvtx3_func_attr__{nvtx3_func_name__}; \
::nvtx3::v1::scoped_range_in<D> const nvtx3_range__{nvtx3_func_attr__};
#else
#define NVTX3_V1_FUNC_RANGE_IN(D) \
roctx_scoped_range_in const roctx_range__{__func__};
#endif
/**
* @brief Convenience macro for generating a range in the specified `domain`
src/include/roctx.h
+165
View File
@@ -0,0 +1,165 @@
/*************************************************************************
* Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
*
* See LICENSE.txt for license information
************************************************************************/
#ifndef RCCL_ROCTX_H
#define RCCL_ROCTX_H
#include <iostream>
#include <string.h>
#include <map>
#include <roctracer/roctx.h>
#include "nvtx3/nvtx3.hpp"
#include "device.h"
#define MAX_MESSAGE_LENGTH 1024
#define NVTX_PAYLOAD_ENTRY_TYPE_REDOP 11 + NVTX_PAYLOAD_ENTRY_TYPE_SCHEMA_ID_STATIC_START
/**
* \brief Equivalent of nvtx types for roctx.
*/
enum roctxPayloadEntryType {
/**
* Only include the required/used types by rccl,
* and needs to be updated in case of new type
* tracing.
*/
ROCTX_PAYLOAD_ENTRY_TYPE_INT,
ROCTX_PAYLOAD_ENTRY_TYPE_SIZE,
ROCTX_PAYLOAD_ENTRY_TYPE_REDOP,
ROCTX_PAYLOAD_NUM_ENTRY_TYPES
};
/**
* \brief Stores the contents of the message that will be used by roctx.
*/
struct roctxPayloadSchemaEntryInfo {
/**
* Description of the data.
*/
const char* name;
/**
* Type of the data.
*/
roctxPayloadEntryType type;
/**
* Union of possible payload types.
*
* Should be in sync with roctxPayloadEntryType.
*/
union {
int typeInt;
size_t typeSize;
ncclDevRedOp_t typeRedOp;
} payload;
};
struct roctxPayloadInfo {
/**
* Payload name. Usually the name of the function/API
* being called from.
*/
const char* id;
/**
* Number of paylod entries
*/
size_t numEntries;
/**
* Pointer to roctxPayloadSchemaEntryInfo elements in memory.
*/
roctxPayloadSchemaEntryInfo* payloadEntries = nullptr;
/**
* Message that will be used by roctx
*/
char* message = nullptr;
};
typedef roctxPayloadInfo* roctxPayloadInfo_t;
extern const char* roctxEntryTypeStr[ROCTX_PAYLOAD_NUM_ENTRY_TYPES];
extern const char* ncclRedOpStr[ncclNumDevRedOps];
/**
* \brief Maps nvtx types to roctx types.
*/
extern std::map<uint64_t, roctxPayloadEntryType> nvtxToRoctx;
/**
* \brief Allocate required memory for roctx
*/
void roctxAlloc(roctxPayloadInfo_t payloadInfo, const size_t numEntries);
/**
* \brief Free all the resources used by roctx
*/
void roctxFree(roctxPayloadInfo_t payloadInfo);
/**
* \brief Extracts payload schema entry info from nvtxPayloadSchemaEntry_t and,
* nvtxPayloadData_t and stores in an array.
*/
void extractPayloadInfo(const nvtxPayloadSchemaEntry_t* schema, const nvtxPayloadData_t* data, const size_t numEntries,
const char* schemaName, roctxPayloadInfo_t payloadInfo);
/**
* \brief Stringify roctxPayloadInfo_t struct. Used as roctx message.
*/
void stringify(roctxPayloadInfo_t payloadInfo);
/**
* \brief Class to make roctx calls scoped.
*/
class roctx_scoped_range_in {
public:
/**
* Construct a 'roctx_scoped_range_in' with specified NVTX params,
* 'numEntries', and 'schemaName'
*/
explicit roctx_scoped_range_in(const nvtxPayloadSchemaEntry_t* schema, const nvtxPayloadData_t* data,
const size_t numEntries, const char* schemaName) noexcept
{
#ifndef ROCTX_NO_IMPL
roctxAlloc(&payloadInfo, numEntries);
extractPayloadInfo(schema, data, numEntries, schemaName, &payloadInfo);
roctxRangePushA(payloadInfo.message);
#endif
}
/**
* Construct a 'roctx_scoped_range_in' with the specified 'message'
*/
explicit roctx_scoped_range_in(const char* message) noexcept
{
#ifndef ROCTX_NO_IMPL
roctxRangePushA(message);
#endif
}
/**
* Default constructor 'roctx_scoped_range_in'
*/
roctx_scoped_range_in() noexcept : roctx_scoped_range_in{""} {/*no impl*/}
/**
* Destroy the roctx_scoped_range_in, ending the ROCTX range event.
*/
~roctx_scoped_range_in() noexcept
{
#ifndef ROCTX_NO_IMPL
roctxRangePop();
roctxFree(&payloadInfo);
#endif
}
private:
roctxPayloadInfo payloadInfo;
};
#endif // RCCL_ROCTX_H
src/misc/roctx.cc
+99
View File
@@ -0,0 +1,99 @@
/*************************************************************************
* Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
*
* See LICENSE.txt for license information
************************************************************************/
#include "roctx.h"
std::map<uint64_t, roctxPayloadEntryType> nvtxToRoctx {
{NVTX_PAYLOAD_ENTRY_TYPE_INT, ROCTX_PAYLOAD_ENTRY_TYPE_INT},
{NVTX_PAYLOAD_ENTRY_TYPE_SIZE, ROCTX_PAYLOAD_ENTRY_TYPE_SIZE},
{NVTX_PAYLOAD_ENTRY_TYPE_REDOP, ROCTX_PAYLOAD_ENTRY_TYPE_REDOP}};
const char* roctxEntryTypeStr[ROCTX_PAYLOAD_NUM_ENTRY_TYPES] = {"ROCTX_PAYLOAD_ENTRY_TYPE_INT", "ROCTX_PAYLOAD_ENTRY_TYPE_SIZE", "ROCTX_PAYLOAD_ENTRY_TYPE_REDOP"};
const char* ncclRedOpStr[ncclNumDevRedOps] = { "Sum", "Prod", "MinMax", "PreMulSum", "SumPostDiv" };
void roctxAlloc(roctxPayloadInfo_t payloadInfo, const size_t numEntries) {
#ifndef ROCTX_NO_IMPL
// Allocate enough memory for numEntries in payloadEntries
payloadInfo->payloadEntries = (roctxPayloadSchemaEntryInfo*)malloc(numEntries * sizeof(roctxPayloadSchemaEntryInfo));
// Allocate memory for the message that will be constructed
payloadInfo->message = (char*)malloc(MAX_MESSAGE_LENGTH * sizeof(char));
#endif
}
void roctxFree(roctxPayloadInfo_t payloadInfo) {
#ifndef ROCTX_NO_IMPL
// Free all the dynamically allocated resources by roctx
if (payloadInfo->payloadEntries) free(payloadInfo->payloadEntries);
if (payloadInfo->message) free((void*)payloadInfo->message);
#endif
}
void extractPayloadInfo(const nvtxPayloadSchemaEntry_t* schema, const nvtxPayloadData_t* data, const size_t numEntries,
const char* schemaName, roctxPayloadInfo_t payloadInfo) {
if (payloadInfo->payloadEntries == nullptr) return;
payloadInfo->id = schemaName;
payloadInfo->numEntries = numEntries;
// Iterate over each entry in the schema
for (size_t i = 0; i < payloadInfo->numEntries; ++i) {
// Populate payload schema entry info for roctx
payloadInfo->payloadEntries[i].name = schema[i].name;
payloadInfo->payloadEntries[i].type = nvtxToRoctx[schema[i].type];
// Offset to index into the data stored in nvtxPayloadData_t->payload
uint64_t offset = schema[i].offset;
const void* entryData = reinterpret_cast<const char*>(data->payload) + offset;
// Populate payload union based on the roctx type
switch (payloadInfo->payloadEntries[i].type) {
case ROCTX_PAYLOAD_ENTRY_TYPE_INT: payloadInfo->payloadEntries[i].payload.typeInt = *reinterpret_cast<const int*>(entryData); break;
case ROCTX_PAYLOAD_ENTRY_TYPE_SIZE: payloadInfo->payloadEntries[i].payload.typeSize = *reinterpret_cast<const size_t*>(entryData); break;
case ROCTX_PAYLOAD_ENTRY_TYPE_REDOP: payloadInfo->payloadEntries[i].payload.typeRedOp = *reinterpret_cast<const ncclDevRedOp_t*>(entryData); break;
default: break;
}
}
// Stringify payloadInfo
stringify(payloadInfo);
}
void stringify(roctxPayloadInfo_t payloadInfo) {
if (!payloadInfo->payloadEntries || !payloadInfo->message) return;
int offset = snprintf(payloadInfo->message, MAX_MESSAGE_LENGTH, "{%s: ", payloadInfo->id);
for (size_t i = 0; i < payloadInfo->numEntries; ++i)
{
roctxPayloadSchemaEntryInfo entry = payloadInfo->payloadEntries[i];
offset += snprintf(payloadInfo->message + offset, MAX_MESSAGE_LENGTH - offset, "%s: ", entry.name);
switch (entry.type)
{
case ROCTX_PAYLOAD_ENTRY_TYPE_INT:
offset += snprintf(payloadInfo->message + offset, MAX_MESSAGE_LENGTH - offset, "%d", entry.payload.typeInt);
break;
case ROCTX_PAYLOAD_ENTRY_TYPE_SIZE:
offset += snprintf(payloadInfo->message + offset, MAX_MESSAGE_LENGTH - offset, "%zu", entry.payload.typeSize);
break;
case ROCTX_PAYLOAD_ENTRY_TYPE_REDOP:
offset += snprintf(payloadInfo->message + offset, MAX_MESSAGE_LENGTH - offset, "%s",
entry.payload.typeRedOp < ncclNumDevRedOps ? ncclRedOpStr[entry.payload.typeRedOp] : "unknown");
break;
default:
offset += snprintf(payloadInfo->message + offset, MAX_MESSAGE_LENGTH - offset, "unknown roctx payload type");
break;
}
if (i != payloadInfo->numEntries - 1)
offset += snprintf(payloadInfo->message + offset, MAX_MESSAGE_LENGTH - offset, ", ");
}
snprintf(payloadInfo->message + offset, MAX_MESSAGE_LENGTH - offset, "}");
}
src/msccl.cc
+12
View File
@@ -44,6 +44,18 @@ ncclResult_t mscclRunAlgo(
void* recvBuff, const size_t recvCounts[], const size_t rDisPls[],
size_t count, ncclDataType_t dataType, int root, int peer, ncclRedOp_t op,
mscclAlgoHandle_t mscclAlgoHandle, ncclComm_t comm, hipStream_t stream) {
struct NvtxParamsMsccl {
size_t sendbytes;
size_t recvbytes;
};
// Just pass the size of one send/recv messages and not the total bytes sent/received.
constexpr nvtxPayloadSchemaEntry_t MscclSchema[] = {
{0, NVTX_PAYLOAD_ENTRY_TYPE_SIZE, "Message size [bytes] (Send)"},
{0, NVTX_PAYLOAD_ENTRY_TYPE_SIZE, "Message size [bytes] (Recv)"}
};
NvtxParamsMsccl payload{sendCounts[comm->rank] * ncclTypeSize(dataType), recvCounts[comm->rank] * ncclTypeSize(dataType)};
NVTX3_FUNC_WITH_PARAMS(MSCCL, MscclSchema, payload)
mscclStatus& status = mscclGetStatus();
struct mscclAlgo* hostAlgo = status.hostAlgos[mscclAlgoHandle];
struct mscclAlgo* devAlgo = status.devAlgos[mscclAlgoHandle];