TensorRT-LLM 0.5.0 源码之二十二

stringUtils

#if ENABLE_BF16#include <cuda_bf16.h>#endif // ENABLE_BF16#include <cuda_fp16.h>

#if ENABLE_BF16static inline std::basic_ostream<char>& operator<<(std::basic_ostream<char>& stream, __nv_bfloat16 const& val){    stream << __bfloat162float(val);    return stream;}#endif // ENABLE_BF16static inline std::basic_ostream<char>& operator<<(std::basic_ostream<char>& stream, __half const& val){    stream << __half2float(val);    return stream;}

std::string fmtstr(char const* format, ...) __attribute__((format(printf, 1, 2)));

template <typename U, typename TStream, typename T>inline TStream& arr2outCasted(TStream& out, T* arr, size_t size){    out << "(";    if (size > 0)    {        for (size_t i = 0; i < size - 1; ++i)        {            out << static_cast<U>(arr[i]) << ", ";        }        out << static_cast<U>(arr[size - 1]);    }    out << ")";    return out;}template <typename TStream, typename T>inline TStream& arr2out(TStream& out, T* arr, size_t size){    return arr2outCasted<T>(out, arr, size);}template <typename T>inline std::string arr2str(T* arr, size_t size){    std::stringstream ss;    return arr2out(ss, arr, size).str();}template <typename T>inline std::string vec2str(std::vector<T> vec){    return arr2str(vec.data(), vec.size());}

inline bool strStartsWith(std::string const& str, std::string const& prefix){    return str.rfind(prefix, 0) == 0;}

std::string vformat(char const* fmt, va_list args){    va_list args0;    va_copy(args0, args);    auto const size = vsnprintf(nullptr, 0, fmt, args0);    if (size <= 0)        return "";    std::string stringBuf(size, char{});    auto const size2 = std::vsnprintf(&stringBuf[0], size + 1, fmt, args);    TLLM_CHECK_WITH_INFO(size2 == size, std::strerror(errno));    return stringBuf;}

std::string vformat(char const* fmt, va_list args){    va_list args0;    va_copy(args0, args);    auto const size = vsnprintf(nullptr, 0, fmt, args0);    if (size <= 0)        return "";    std::string stringBuf(size, char{});    auto const size2 = std::vsnprintf(&stringBuf[0], size + 1, fmt, args);    TLLM_CHECK_WITH_INFO(size2 == size, std::strerror(errno));    return stringBuf;}

stlUtil

template <typename TInputIt, typename TOutputIt, typename TBinOp>constexpr TOutputIt basicInclusiveScan(TInputIt first, TInputIt last, TOutputIt dFirst, TBinOp op){    if (first != last)    {        auto val = *first;        while (true)        {            *dFirst = val;            ++dFirst;            ++first;            if (first == last)            {                break;            }            val = op(std::move(val), *first);        }    }    return dFirst;}template <typename TInputIt, typename TOutputIt>constexpr TOutputIt inclusiveScan(TInputIt first, TInputIt last, TOutputIt dFirst){#if defined(__GNUC__) && __GNUC__ <= 8    return basicInclusiveScan(first, last, dFirst, std::plus<>{});#else    return std::inclusive_scan(first, last, dFirst);#endif}

InclusiveScan:input: 0, 1, 2, 3output: 0, 1, 3, 6

template <typename TInputIt, typename TOutputIt, typename T, typename TBinOp>constexpr TOutputIt basicExclusiveScan(TInputIt first, TInputIt last, TOutputIt dFirst, T init, TBinOp op){    if (first != last)    {        while (true)        {            T tmp{op(init, *first)};            *dFirst = init;            ++dFirst;            ++first;            if (first == last)            {                break;            }            init = std::move(tmp);        }    }    return dFirst;}template <typename TInputIt, typename TOutputIt, typename T>constexpr TOutputIt exclusiveScan(TInputIt first, TInputIt last, TOutputIt dFirst, T init){#if defined(__GNUC__) && __GNUC__ <= 8    return basicExclusiveScan(first, last, dFirst, std::move(init), std::plus<>{});#else    return std::exclusive_scan(first, last, dFirst, std::move(init));#endif}

EclusiveScan：Input: 0, 1, 2, 3, init=5Output: 5, 5, 6, 8

独占前缀和是指，输出序列中的每个元素是输入序列中该位置之前所有元素（不包括当前位置元素）经过累积操作的结果。

用一个简单例子说明。假设二元操作 op是加法 +，初始值 init为 0：输入序列: [a, b, c, d]输出序列: [0, a, a+b, a+b+c]（注意输出序列比输入序列多一个元素，但此函数输出序列长度与输入相同）

cublasVersionCheck

#define TLLM_CUBLAS_VERSION_CALC(MAJOR, MINOR, PATCH) (MAJOR * 10000 + MINOR * 100 + PATCH)#define TLLM_CUBLAS_VER_LE(MAJOR, MINOR, PATCH)                                                                        \    TLLM_CUBLAS_VERSION_CALC(CUBLAS_VER_MAJOR, CUBLAS_VER_MINOR, CUBLAS_VER_PATCH)                                     \    <= TLLM_CUBLAS_VERSION_CALC(MAJOR, MINOR, PATCH)#define TLLM_CUBLAS_VER_LT(MAJOR, MINOR, PATCH)                                                                        \    TLLM_CUBLAS_VERSION_CALC(CUBLAS_VER_MAJOR, CUBLAS_VER_MINOR, CUBLAS_VER_PATCH)                                     \        < TLLM_CUBLAS_VERSION_CALC(MAJOR, MINOR, PATCH)#define TLLM_CUBLAS_VER_GE(MAJOR, MINOR, PATCH)                                                                        \    TLLM_CUBLAS_VERSION_CALC(CUBLAS_VER_MAJOR, CUBLAS_VER_MINOR, CUBLAS_VER_PATCH)                                     \    >= TLLM_CUBLAS_VERSION_CALC(MAJOR, MINOR, PATCH)#define TLLM_CUBLAS_VER_GT(MAJOR, MINOR, PATCH)                                                                        \    TLLM_CUBLAS_VERSION_CALC(CUBLAS_VER_MAJOR, CUBLAS_VER_MINOR, CUBLAS_VER_PATCH)                                     \        > TLLM_CUBLAS_VERSION_CALC(MAJOR, MINOR, PATCH)

nvtxUtils

#include <nvtx3/nvtx3.hpp>#include <array>namespace tensorrt_llm::common::nvtx{inline nvtx3::color nextColor(){#if !defined(NVTX_DISABLE)    constexpr std::array kColors{nvtx3::color{0xff00ff00}, nvtx3::color{0xff0000ff}, nvtx3::color{0xffffff00},        nvtx3::color{0xffff00ff}, nvtx3::color{0xff00ffff}, nvtx3::color{0xffff0000}, nvtx3::color{0xffffffff}};    constexpr auto numColors = kColors.size();    static thread_local std::size_t colorId = 0;    auto const color = kColors[colorId];    colorId = colorId + 1 >= numColors ? 0 : colorId + 1;    return color;#else    return nvtx3::color{0};#endif}} // namespace tensorrt_llm::common::nvtx#define NVTX3_SCOPED_RANGE(range) ::nvtx3::scoped_range range##_range(::tensorrt_llm::common::nvtx::nextColor(), #range)

这段代码是用于GPU性能分析的工具，它利用NVIDIA的NVTX（NVIDIA Tools Extension）库来标记代码段，以便在性能分析工具（如Nsight Systems、Nsight Compute）中可视化程序的执行流程和时间分布。

下面逐部分解释代码的功能：

头文件与命名空间

#include <nvtx3/nvtx3.hpp>

• nvtx3/nvtx3.hpp：引入了NVTX v3库的功能，用于创建范围标记和颜色定义。

nextColor() 函数：循环分配颜色

inline nvtx3::color nextColor(){#if !defined(NVTX_DISABLE)    constexpr std::array kColors{nvtx3::color{0xff00ff00}, nvtx3::color{0xff0000ff},                                 nvtx3::color{0xffffff00}, nvtx3::color{0xffff00ff},                                 nvtx3::color{0xff00ffff}, nvtx3::color{0xffff0000},                                 nvtx3::color{0xffffffff}};    constexpr auto numColors = kColors.size();    static thread_local std::size_t colorId = 0;    auto const color = kColors[colorId];    colorId = colorId + 1 >= numColors ? 0 : colorId + 1;    return color;#else    return nvtx3::color{0}; // 如果禁用NVTX，返回黑色（无操作）#endif}

• 功能：该函数从一个预定义的七种颜色数组中循环返回颜色。颜色以ARGB格式表示（例如0xff00ff00是绿色）。
• 线程安全：使用static thread_local确保每个线程有独立的颜色计数器，避免多线程竞争。
• 条件编译：如果定义了NVTX_DISABLE宏，则返回黑色（效果上相当于禁用NVTX标记），这提供了灵活的编译时控制。

NVTX3_SCOPED_RANGE 宏：创建作用域标记

#define NVTX3_SCOPED_RANGE(range) ::nvtx3::scoped_range range##_range(::tensorrt_llm::common::nvtx::nextColor(), #range)

• 功能：这是一个宏，用于创建一个作用域范围的NVTX标记。当程序执行进入该宏定义的代码块时，性能分析工具的时间线上会开始一个具有指定颜色和名称的区间；离开作用域时，区间自动结束。
• 参数说明：

• range：用户指定的范围名称（例如NVTX3_SCOPED_RANGE(forward_pass)会生成名为"forward_pass"的区间）。
• range##_range：生成一个唯一的变量名，确保在同一作用域内可多次使用该宏。
• nextColor()：为每个范围分配一个循环颜色，帮助在时间线上直观区分不同代码段。

代码应用场景与示例

这段代码通常用于GPU加速的应用程序（如深度学习框架TensorRT-LLM）中，进行精细化的性能分析。例如：

void run_mlp(){    NVTX3_SCOPED_RANGE(forward_pass); // 标记前向传播范围    // ... 前向传播计算 ...    NVTX3_SCOPED_RANGE(backward_pass); // 标记反向传播范围    // ... 反向传播计算 ...}