TensorDataType
typedef enum datatype_enum
{
TYPE_INVALID,
TYPE_BOOL,
TYPE_UINT8,
TYPE_UINT16,
TYPE_UINT32,
TYPE_UINT64,
TYPE_INT8,
TYPE_INT16,
TYPE_INT32,
TYPE_INT64,
TYPE_FP16,
TYPE_FP32,
TYPE_FP64,
TYPE_BYTES,
TYPE_BF16,
TYPE_FP8_E4M3,
TYPE_STR,
TYPE_VOID,
TYPE_INT32_PTR,
} DataType;
template <typename T>
struct TensorDataType
{
};
template <>
struct TensorDataType<bool>
{
static constexpr DataType value = TYPE_BOOL;
};
template <>
struct TensorDataType<std::uint8_t>
{
static constexpr DataType value = TYPE_UINT8;
};
template <>
struct TensorDataType<std::uint16_t>
{
static constexpr DataType value = TYPE_UINT16;
};
template <>
struct TensorDataType<std::uint32_t>
{
static constexpr DataType value = TYPE_UINT32;
};
template <>
struct TensorDataType<std::uint64_t>
{
static constexpr DataType value = TYPE_UINT64;
};
#if !defined(_WIN32)
template <>
struct TensorDataType<unsigned long long>
{
static constexpr DataType value = TYPE_UINT64;
};
#endif // !defined(_WIN32)
static_assert(sizeof(std::uint64_t) == sizeof(unsigned long long), "");
template <>
struct TensorDataType<std::int8_t>
{
static constexpr DataType value = TYPE_INT8;
};
template <>
struct TensorDataType<std::int16_t>
{
static constexpr DataType value = TYPE_INT16;
};
template <>
struct TensorDataType<std::int32_t>
{
static constexpr DataType value = TYPE_INT32;
};
template <>
struct TensorDataType<std::int64_t>
{
static constexpr DataType value = TYPE_INT64;
};
template <>
struct TensorDataType<half>
{
static constexpr DataType value = TYPE_FP16;
};
template <>
struct TensorDataType<float>
{
static constexpr DataType value = TYPE_FP32;
};
template <>
struct TensorDataType<double>
{
static constexpr DataType value = TYPE_FP64;
};
template <>
struct TensorDataType<char>
{
static constexpr DataType value = TYPE_BYTES;
};
#ifdef ENABLE_BF16
template <>
struct TensorDataType<__nv_bfloat16>
{
static constexpr DataType value = TYPE_BF16;
};
#endif
#ifdef ENABLE_FP8
template <>
struct TensorDataType<__nv_fp8_e4m3>
{
static constexpr DataType value = TYPE_FP8_E4M3;
};
#endif
template <>
struct TensorDataType<std::string>
{
static constexpr DataType value = TYPE_STR;
};
template <>
struct TensorDataType<void>
{
static constexpr DataType value = TYPE_VOID;
};
template <>
struct TensorDataType<int*>
{
static constexpr DataType value = TYPE_INT32_PTR;
};
template <>
struct TensorDataType<const int*>
{
static constexpr DataType value = TYPE_INT32_PTR;
};
template <typename T>
DataType getTensorType()
{
return TensorDataType<typename std::remove_cv<T>::type>::value;
}
Tensor
typedef enum memorytype_enum
{
MEMORY_CPU,
MEMORY_CPU_PINNED,
MEMORY_GPU
} MemoryType;
class Tensor
{
public:
// Do not write to these variables directly. Use copy / move constructors instead.
MemoryType where;
DataType type;
std::vector<size_t> shape;
void const* data; // TODO modify from const void* to void* const
Tensor();
Tensor(MemoryType _where, DataType _type, std::vector<size_t> const& _shape, void const* _data);
std::size_t size() const;
std::size_t sizeBytes() const;
std::string whereToString() const;
std::string toString() const;
std::string getNumpyTypeDesc(DataType type) const;
void saveNpy(const std::string& filename) const;
static Tensor loadNpy(const std::string& npy_file, const MemoryType where);
static DataType typeFromNumpyDesc(std::string type);
static size_t getTypeSize(DataType type);
template <typename T>
inline T getVal(size_t index) const
{
TLLM_LOG_DEBUG("%s start", __PRETTY_FUNCTION__);
TLLM_CHECK(where == MEMORY_CPU);
TLLM_CHECK(data != nullptr);
TLLM_CHECK_WITH_INFO(index < size(), "index is larger than buffer size");
if (getTensorType<T>() != type)
{
TLLM_LOG_DEBUG("getVal with type %s, but data type is: %s", getNumpyTypeDesc(getTensorType<T>()).c_str(),
getNumpyTypeDesc(type).c_str());
}
return ((T*) data)[index];
}
template <typename T>
inline T getVal() const
{
TLLM_LOG_DEBUG("%s start", __PRETTY_FUNCTION__);
if (getTensorType<T>() != type)
{
TLLM_LOG_DEBUG("getVal with type %s, but data type is: %s", getNumpyTypeDesc(getTensorType<T>()).c_str(),
getNumpyTypeDesc(type).c_str());
}
return getVal<T>(0);
}
template <typename T>
inline T* getPtr() const
{
TLLM_LOG_DEBUG("%s start", __PRETTY_FUNCTION__);
if (getTensorType<T>() != type)
{
TLLM_LOG_DEBUG("getPtr with type %s, but data type is: %s", getNumpyTypeDesc(getTensorType<T>()).c_str(),
getNumpyTypeDesc(type).c_str());
}
return (T*) data;
}
inline void* getPtrWithOffset(size_t offset) const
{
TLLM_LOG_DEBUG("%s start", __PRETTY_FUNCTION__);
if (data == nullptr)
{
return (void*) data;
}
else
{
TLLM_CHECK_WITH_INFO(offset < size(), "offset is larger than buffer size");
return (void*) ((char*) data + offset * Tensor::getTypeSize(type));
}
}
template <typename T>
inline T* getPtrWithOffset(size_t offset) const
{
TLLM_LOG_DEBUG("%s start", __PRETTY_FUNCTION__);
if (getTensorType<T>() != type)
{
TLLM_LOG_DEBUG("getVal with type %s, but data type is: %s", getNumpyTypeDesc(getTensorType<T>()).c_str(),
getNumpyTypeDesc(type).c_str());
}
if (data == nullptr)
{
return (T*) data;
}
else
{
TLLM_CHECK_WITH_INFO(
offset < size(), fmtstr("offset (%lu) is larger than buffer size (%lu)", offset, size()));
return ((T*) data) + offset;
}
}
template <typename T>
T max() const
{
if (getTensorType<T>() != type)
{
TLLM_LOG_DEBUG("getVal with type %s, but data type is: %s", getNumpyTypeDesc(getTensorType<T>()).c_str(),
getNumpyTypeDesc(type).c_str());
}
TLLM_CHECK_WITH_INFO(shape.size() > 0 && data != nullptr, "Should be a non-empty tensor.");
TLLM_CHECK_WITH_INFO(where == MEMORY_CPU || where == MEMORY_CPU_PINNED,
"max() supports MEMORY_CPU or MEMORY_CPU_PINNED tensor.");
size_t max_idx = 0;
T max_val = getVal<T>(max_idx);
for (size_t i = 1; i < size(); ++i)
{
T val = getVal<T>(i);
if (val > max_val)
{
max_idx = i;
max_val = val;
}
}
return max_val;
}
template <typename T>
T min() const
{
if (getTensorType<T>() != type)
{
TLLM_LOG_DEBUG("getVal with type %s, but data type is: %s", getNumpyTypeDesc(getTensorType<T>()).c_str(),
getNumpyTypeDesc(type).c_str());
}
TLLM_CHECK_WITH_INFO(shape.size() > 0 && data != nullptr, "Should be a non-empty tensor.");
TLLM_CHECK_WITH_INFO(where == MEMORY_CPU || where == MEMORY_CPU_PINNED,
"min() supports MEMORY_CPU or MEMORY_CPU_PINNED tensor.");
size_t min_idx = 0;
T min_val = getVal<T>(min_idx);
for (size_t i = 1; i < size(); ++i)
{
T val = getVal<T>(i);
if (val < min_val)
{
min_idx = i;
min_val = val;
}
}
return min_val;
}
template <typename T>
T any(T val) const
{
if (getTensorType<T>() != type)
{
TLLM_LOG_DEBUG("getVal with type %s, but data type is: %s", getNumpyTypeDesc(getTensorType<T>()).c_str(),
getNumpyTypeDesc(type).c_str());
}
TLLM_CHECK_WITH_INFO(shape.size() > 0 && data != nullptr, "Should be a non-empty tensor.");
TLLM_CHECK_WITH_INFO(where == MEMORY_CPU || where == MEMORY_CPU_PINNED,
"any() supports MEMORY_CPU or MEMORY_CPU_PINNED tensor.");
for (size_t i = 0; i < size(); ++i)
{
if (getVal<T>(i) == val)
{
return true;
}
}
return false;
}
template <typename T>
T all(T val) const
{
if (getTensorType<T>() != type)
{
TLLM_LOG_DEBUG("getVal with type %s, but data type is: %s", getNumpyTypeDesc(getTensorType<T>()).c_str(),
getNumpyTypeDesc(type).c_str());
}
TLLM_CHECK_WITH_INFO(shape.size() > 0 && data != nullptr, "Should be a non-empty tensor.");
TLLM_CHECK_WITH_INFO(where == MEMORY_CPU || where == MEMORY_CPU_PINNED,
"all() supports MEMORY_CPU or MEMORY_CPU_PINNED tensor.");
for (size_t i = 0; i < size(); ++i)
{
if (getVal<T>(i) != val)
{
return false;
}
}
return true;
}
void updateShape(size_t idx, size_t val)
{
// TODO: find a better way to update the shape
std::vector<size_t>& shape_ref = const_cast<std::vector<size_t>&>(shape);
shape_ref[idx] = val;
}
inline bool isValid() const
{
return size() > 0 && data != nullptr;
}
Tensor slice(std::vector<size_t> shape, size_t offset = 0) const;
private:
static void parseNpyIntro(FILE*& f_ptr, uint32_t& header_len, uint32_t& start_data);
static int parseNpyHeader(FILE*& f_ptr, uint32_t header_len, DataType& type, std::vector<size_t>& shape);
};
Tensor::Tensor()
: // a none tensor.
where(MEMORY_CPU)
, type(TYPE_INVALID)
, shape({})
, data(nullptr)
{
}
Tensor::Tensor(MemoryType _where, DataType _type, std::vector<size_t> const& _shape, void const* _data)
: where(_where)
, type(_type)
, shape(_shape)
, data(_data)
{
}
void Tensor::parseNpyIntro(FILE*& f_ptr, uint32_t& header_len, uint32_t& start_data)
{
const char magic[]
= "\x93"
"NUMPY";
char magic_test[sizeof(magic)] = "\0";
size_t n_elems = fread((void*) magic_test, sizeof(char), sizeof(magic) - 1, f_ptr);
if (n_elems != sizeof(magic) - 1 || std::string(magic) != std::string(magic_test))
{
throw std::runtime_error("Could read magic token in NPY file");
}
uint8_t npy_major = 0;
uint8_t npy_minor = 0;
n_elems = fread((void*) &npy_major, sizeof(uint8_t), 1, f_ptr);
n_elems += fread((void*) &npy_minor, sizeof(uint8_t), 1, f_ptr);
if (npy_major == 1)
{
uint16_t header_len_u16 = 0;
n_elems = fread((void*) &header_len_u16, sizeof(uint16_t), 1, f_ptr);
header_len = header_len_u16;
}
else if (npy_major == 2)
{
uint32_t header_len_u32 = 0;
n_elems = fread((void*) &header_len_u32, sizeof(uint32_t), 1, f_ptr);
header_len = header_len_u32;
}
else
{
throw std::runtime_error("Unsupported npy version: " + std::to_string(npy_major));
}
start_data = 8 + 2 * npy_major + header_len;
}
int Tensor::parseNpyHeader(FILE*& f_ptr, uint32_t header_len, DataType& type, std::vector<size_t>& shape)
{
char* header_c = (char*) malloc(header_len * sizeof(char));
size_t n_elems = fread((void*) header_c, sizeof(char), header_len, f_ptr);
if (n_elems != header_len)
{
free(header_c);
return -1;
}
std::string header(header_c, header_len);
free(header_c);
size_t start, end;
start = header.find("'descr'") + 7;
start = header.find("'", start);
end = header.find("'", start + 1);
type = typeFromNumpyDesc(header.substr(start + 2, end - start - 2));
start = header.find("'fortran_order'") + 15;
start = header.find(":", start);
end = header.find(",", start + 1);
if (header.substr(start + 1, end - start - 1).find("False") == std::string::npos)
{
throw std::runtime_error("Unsupported value for fortran_order while reading npy file");
}
start = header.find("'shape'") + 7;
start = header.find("(", start);
end = header.find(")", start + 1);
std::istringstream shape_stream(header.substr(start + 1, end - start - 1));
std::string token;
shape.clear();
while (std::getline(shape_stream, token, ','))
{
if (token.find_first_not_of(' ') == std::string::npos)
{
break;
}
shape.push_back(std::stoul(token));
}
return 0;
}
Tensor Tensor::loadNpy(const std::string& npy_file, const MemoryType where)
{
DataType type;
std::vector<size_t> shape;
FILE* f_ptr = fopen(npy_file.c_str(), "rb");
if (f_ptr == nullptr)
{
throw std::runtime_error("Could not open file " + npy_file);
}
uint32_t header_len, start_data;
parseNpyIntro(f_ptr, header_len, start_data);
parseNpyHeader(f_ptr, header_len, type, shape);
const size_t size = std::accumulate(shape.begin(), shape.end(), size_t{1}, std::multiplies<size_t>());
void* data_cpu = malloc(size * Tensor::getTypeSize(type));
void* data = data_cpu;
size_t n_elems = fread(data_cpu, Tensor::getTypeSize(type), size, f_ptr);
TLLM_CHECK_WITH_INFO(n_elems == size, "reading tensor failed");
if (where == MEMORY_GPU)
{
cudaMalloc(&data, size * Tensor::getTypeSize(type));
cudaMemcpy(data, data_cpu, size * Tensor::getTypeSize(type), cudaMemcpyHostToDevice);
free(data_cpu);
}
fclose(f_ptr);
return Tensor(where, type, shape, data);
}
size_t Tensor::size() const
{
if (data == nullptr || shape.size() == 0)
{
return 0;
}
return std::accumulate(shape.begin(), shape.end(), (size_t) 1, std::multiplies<size_t>());
}
size_t Tensor::sizeBytes() const
{
return size() * Tensor::getTypeSize(type);
}
std::string Tensor::whereToString() const
{
static const std::unordered_map<MemoryType, std::string> mem_to_string{
{MEMORY_CPU, "CPU"}, {MEMORY_CPU_PINNED, "CPU_PINNED"}, {MEMORY_GPU, "GPU"}};
return mem_to_string.at(where);
}
std::string Tensor::toString() const
{
std::string memtype_str = whereToString();
static const std::unordered_map<DataType, std::string> type_to_string{
{TYPE_BOOL, "BOOL"},
{TYPE_UINT8, "UINT8"},
{TYPE_UINT16, "UINT16"},
{TYPE_UINT32, "UINT32"},
{TYPE_UINT64, "UINT64"},
{TYPE_INT8, "INT8"},
{TYPE_INT16, "INT16"},
{TYPE_INT32, "INT32"},
{TYPE_INT64, "INT64"},
{TYPE_BF16, "BF16"},
{TYPE_FP16, "FP16"},
{TYPE_FP32, "FP32"},
{TYPE_FP64, "FP64"},
{TYPE_BYTES, "BYTES"},
{TYPE_INVALID, "INVALID"},
{TYPE_FP8_E4M3, "E4M3"},
{TYPE_VOID, "VOID"},
};
return fmtstr("Tensor[where=%s, type=%s, shape=%s, data=%p]", memtype_str.c_str(), type_to_string.at(type).c_str(),
vec2str(shape).c_str(), data);
}
DataType Tensor::typeFromNumpyDesc(std::string type)
{
static const std::unordered_map<std::string, DataType> type_map{{"?", TYPE_BOOL}, {"b", TYPE_BYTES},
{"u1", TYPE_UINT8}, {"u2", TYPE_UINT16}, {"u4", TYPE_UINT32}, {"u8", TYPE_UINT64}, {"i1", TYPE_INT8},
{"i2", TYPE_INT16}, {"i4", TYPE_INT32}, {"i8", TYPE_INT64}, {"f2", TYPE_FP16}, {"f4", TYPE_FP32},
{"f8", TYPE_FP64}};
TLLM_CHECK_WITH_INFO(type_map.count(type) > 0, "numpy data type '" + type + "' not supported");
return type_map.at(type);
}
size_t Tensor::getTypeSize(DataType type)
{
static const std::unordered_map<DataType, size_t> type_map{{TYPE_BOOL, sizeof(bool)}, {TYPE_BYTES, sizeof(char)},
{TYPE_UINT8, sizeof(uint8_t)}, {TYPE_UINT16, sizeof(uint16_t)}, {TYPE_UINT32, sizeof(uint32_t)},
{TYPE_UINT64, sizeof(uint64_t)}, {TYPE_INT8, sizeof(int8_t)}, {TYPE_INT16, sizeof(int16_t)},
{TYPE_INT32, sizeof(int32_t)}, {TYPE_INT64, sizeof(int64_t)},
#ifdef ENABLE_BF16
{TYPE_BF16, sizeof(__nv_bfloat16)},
#endif
#ifdef ENABLE_FP8
{TYPE_FP8_E4M3, sizeof(__nv_fp8_e4m3)},
#endif
{TYPE_FP16, sizeof(half)}, {TYPE_FP32, sizeof(float)}, {TYPE_FP64, sizeof(double)}};
return type_map.at(type);
}
std::string Tensor::getNumpyTypeDesc(DataType type) const
{
static const std::unordered_map<DataType, std::string> type_map{{TYPE_INVALID, "x"}, {TYPE_BOOL, "?"},
{TYPE_BYTES, "b"}, {TYPE_UINT8, "u1"}, {TYPE_UINT16, "u2"}, {TYPE_UINT32, "u4"}, {TYPE_UINT64, "u8"},
{TYPE_INT8, "i1"}, {TYPE_INT16, "i2"}, {TYPE_INT32, "i4"}, {TYPE_INT64, "i8"}, {TYPE_FP16, "f2"},
{TYPE_FP32, "f4"}, {TYPE_FP64, "f8"}};
if (type == TYPE_BF16)
{
TLLM_LOG_WARNING(
"getNumpyTypeDesc(TYPE_BF16) returns an invalid type 'x' since Numpy doesn't "
"support bfloat16 as of now, it will be properly extended if numpy supports. "
"Please refer for the discussions https://github.com/numpy/numpy/issues/19808.");
}
return type_map.count(type) > 0 ? type_map.at(type) : "x";
}
void Tensor::saveNpy(const std::string& filename) const
{
// Save tensor to NPY 1.0 format (see https://numpy.org/neps/nep-0001-npy-format.html)
void* cpu_data = (void*) data;
bool is_data_temp = false;
size_t tensor_size = size();
#ifdef ENABLE_BF16
if (type == TYPE_BF16)
{
TLLM_CHECK(where == MemoryType::MEMORY_GPU);
float* data_fp32 = nullptr;
cudaMalloc(&data_fp32, tensor_size * sizeof(float));
invokeCudaD2DcpyConvert(data_fp32, static_cast<const __nv_bfloat16*>(data), tensor_size);
Tensor{where, TYPE_FP32, shape, data_fp32}.saveNpy(filename);
cudaFree(data_fp32);
return;
}
#endif
if (where == MemoryType::MEMORY_GPU)
{
cpu_data = malloc(tensor_size * Tensor::getTypeSize(type));
is_data_temp = true;
cudaDeviceSynchronize();
cudaMemcpy(cpu_data, data, tensor_size * Tensor::getTypeSize(type), cudaMemcpyDeviceToHost);
}
const char magic[]
= "\x93"
"NUMPY";
const uint8_t npy_major = 1;
const uint8_t npy_minor = 0;
std::stringstream header_stream;
header_stream << "{'descr': '" << getNumpyTypeDesc(type) << "', 'fortran_order': False, 'shape': (";
for (size_t i = 0; i < shape.size(); ++i)
{
header_stream << shape[i];
if (i + 1 < shape.size() || shape.size() == 1)
{
header_stream << ", ";
}
}
header_stream << ")}";
int base_length = 6 + 4 + header_stream.str().size();
int pad_length = 16 * ((base_length + 1 + 15) / 16); // Take ceiling of base_length + 1 (for '\n' ending)
for (int i = 0; i < pad_length - base_length; ++i)
{
header_stream << ((i == pad_length - base_length - 1) ? "\n" : "\x20");
}
std::string header = header_stream.str();
const uint16_t header_len = header.size();
FILE* f_ptr = fopen(filename.c_str(), "wb");
TLLM_CHECK_WITH_INFO(f_ptr != nullptr, fmtstr("Unable to open %s for writing.\n", filename.c_str()));
fwrite(magic, sizeof(char), sizeof(magic) - 1, f_ptr);
fwrite(&npy_major, sizeof(uint8_t), 1, f_ptr);
fwrite(&npy_minor, sizeof(uint8_t), 1, f_ptr);
fwrite(&header_len, sizeof(uint16_t), 1, f_ptr);
fwrite(header.c_str(), sizeof(char), header_len, f_ptr);
fwrite(cpu_data, Tensor::getTypeSize(type), tensor_size, f_ptr);
fclose(f_ptr);
if (is_data_temp)
{
free(cpu_data);
}
}
Tensor Tensor::slice(std::vector<size_t> shape, size_t offset) const
{
if (this->data != nullptr)
{
size_t n_elts = this->size();
size_t n_sliced_elts = std::accumulate(shape.begin(), shape.end(), size_t{1}, std::multiplies<size_t>());
TLLM_CHECK_WITH_INFO(n_sliced_elts + offset <= n_elts,
fmtstr("The number (%ld) of elements of sliced tensor exceeds that (%ld) of the original tensor",
n_sliced_elts + offset, n_elts));
}
return Tensor(this->where, this->type, shape, this->getPtrWithOffset(offset));
}
TensorMap
class TensorMap
{
private:
std::unordered_map<std::string, Tensor> tensor_map_;
public:
TensorMap() = default;
TensorMap(const std::unordered_map<std::string, Tensor>& tensor_map);
TensorMap(const std::vector<Tensor>& tensor_map);
TensorMap(std::initializer_list<std::pair<std::string, Tensor>> tensor_map);
~TensorMap();
inline size_t size() const
{
return tensor_map_.size();
}
inline bool contains(const std::string& key) const
{
TLLM_LOG_DEBUG("%s for key: %s", __PRETTY_FUNCTION__, key.c_str());
return tensor_map_.find(key) != tensor_map_.end();
}
std::vector<std::string> keys() const;
inline void insert(const std::string& key, const Tensor& value)
{
TLLM_CHECK_WITH_INFO(!contains(key), fmtstr("Duplicated key %s", key.c_str()));
TLLM_CHECK_WITH_INFO(
value.isValid(), fmtstr("A none tensor or nullptr is not allowed (key is %s)", key.c_str()));
tensor_map_.insert({key, value});
}
inline void insertIfValid(const std::string& key, const Tensor& value)
{
if (value.isValid())
{
insert({key, value});
}
}
inline void insert(std::pair<std::string, Tensor> p)
{
tensor_map_.insert(p);
}
// prevent converting int or size_t to string automatically
Tensor at(int tmp)= delete;
Tensor at(size_t tmp)= delete;
inline Tensor& at(const std::string& key)
{
TLLM_LOG_DEBUG("%s for key %s", __PRETTY_FUNCTION__, key.c_str());
TLLM_CHECK_WITH_INFO(contains(key),
fmtstr(
"Cannot find a tensor of name %s in the tensor map (keys: %s)", key.c_str(), vec2str(keys()).c_str()));
return tensor_map_.at(key);
}
inline Tensor at(const std::string& key) const
{
TLLM_CHECK_WITH_INFO(contains(key),
fmtstr(
"Cannot find a tensor of name %s in the tensor map (keys: %s)", key.c_str(), vec2str(keys()).c_str()));
return tensor_map_.at(key);
}
inline std::optional<Tensor> atOpt(const std::string& key) const
{
if (contains(key))
return tensor_map_.at(key);
else
return std::nullopt;
}
inline Tensor& at(const std::string& key, Tensor& default_tensor)
{
TLLM_LOG_DEBUG("%s for key %s", __PRETTY_FUNCTION__, key.c_str());
if (contains(key))
{
return tensor_map_.at(key);
}
return default_tensor;
}
inline Tensor at(const std::string& key, Tensor& default_tensor) const
{
TLLM_LOG_DEBUG("%s for key %s", __PRETTY_FUNCTION__, key.c_str());
if (contains(key))
{
return tensor_map_.at(key);
}
return default_tensor;
}
inline Tensor& at(const std::string& key, Tensor&& default_tensor)
{
TLLM_LOG_DEBUG("%s for key %s", __PRETTY_FUNCTION__, key.c_str());
if (contains(key))
{
return tensor_map_.at(key);
}
return default_tensor;
}
inline Tensor at(const std::string& key, Tensor&& default_tensor) const
{
if (contains(key))
{
return tensor_map_.at(key);
}
return default_tensor;
}
template <typename T>
inline T getVal(const std::string& key) const
{
TLLM_CHECK_WITH_INFO(contains(key),
fmtstr(
"Cannot find a tensor of name %s in the tensor map (keys: %s)", key.c_str(), vec2str(keys()).c_str()));
return tensor_map_.at(key).getVal<T>();
}
template <typename T>
inline std::optional<T> getValOpt(const std::string& key) const
{
if (contains(key))
{
return tensor_map_.at(key).getVal<T>();
}
else
{
return std::nullopt;
}
}
template <typename T>
inline T getVal(const std::string& key, T default_value) const
{
if (contains(key))
{
return tensor_map_.at(key).getVal<T>();
}
return default_value;
}
template <typename T>
inline T getValWithOffset(const std::string& key, size_t index) const
{
TLLM_CHECK_WITH_INFO(contains(key),
fmtstr(
"Cannot find a tensor of name %s in the tensor map (keys: %s)", key.c_str(), vec2str(keys()).c_str()));
return tensor_map_.at(key).getVal<T>(index);
}
template <typename T>
inline T getValWithOffset(const std::string& key, size_t index, T default_value) const
{
if (contains(key))
{
return tensor_map_.at(key).getVal<T>(index);
}
return default_value;
}
template <typename T>
inline T* getPtr(const std::string& key) const
{
TLLM_CHECK_WITH_INFO(contains(key),
fmtstr(
"Cannot find a tensor of name %s in the tensor map (keys: %s)", key.c_str(), vec2str(keys()).c_str()));
return tensor_map_.at(key).getPtr<T>();
}
template <typename T>
inline T* getPtr(const std::string& key, T* default_ptr) const
{
if (contains(key))
{
return tensor_map_.at(key).getPtr<T>();
}
return default_ptr;
}
template <typename T>
inline T* getPtrWithOffset(const std::string& key, size_t index) const
{
TLLM_CHECK_WITH_INFO(contains(key),
fmtstr(
"Cannot find a tensor of name %s in the tensor map (keys: %s)", key.c_str(), vec2str(keys()).c_str()));
return tensor_map_.at(key).getPtrWithOffset<T>(index);
}
template <typename T>
inline T* getPtrWithOffset(const std::string& key, size_t index, T* default_ptr) const
{
if (contains(key))
{
return tensor_map_.at(key).getPtrWithOffset<T>(index);
}
return default_ptr;
}
inline std::unordered_map<std::string, Tensor> getMap() const
{
return tensor_map_;
}
inline std::unordered_map<std::string, Tensor>::iterator begin()
{
return tensor_map_.begin();
}
inline std::unordered_map<std::string, Tensor>::iterator end()
{
return tensor_map_.end();
}
std::string toString();
static TensorMap fromNpyFolder(const std::string& base_folder);
void saveNpy(const std::string& base_folder);
};
TensorMap::TensorMap(const std::unordered_map<std::string, Tensor>& tensor_map)
{
for (auto& kv : tensor_map)
{
if (kv.second.isValid())
{
insert(kv.first, kv.second);
}
else
{
TLLM_LOG_DEBUG(fmtstr("%s is not a valid tensor, skipping insert into TensorMap", kv.first.c_str()));
}
}
}
TensorMap::TensorMap(const std::vector<Tensor>& tensor_map)
{
for (size_t i = 0; i < tensor_map.size(); i++)
{
insert(std::to_string(i), tensor_map[i]);
}
}
TensorMap::TensorMap(std::initializer_list<std::pair<std::string, Tensor>> tensor_map)
{
for (auto& pair : tensor_map)
{
if (pair.second.isValid())
{
insert(pair.first, pair.second);
}
else
{
TLLM_LOG_DEBUG(fmtstr("%s is not a valid tensor, skipping insert into TensorMap", pair.first.c_str()));
}
}
}
TensorMap::~TensorMap()
{
tensor_map_.clear();
}
std::vector<std::string> TensorMap::keys() const
{
std::vector<std::string> key_names;
for (auto& kv : tensor_map_)
{
key_names.push_back(kv.first);
}
return key_names;
}
std::string TensorMap::toString()
{
std::stringstream ss;
ss << "{";
std::vector<std::string> key_names = keys();
for (size_t i = 0; i < tensor_map_.size(); ++i)
{
ss << key_names[i] << ": " << at(key_names[i]).toString();
if (i < tensor_map_.size() - 1)
{
ss << ", ";
}
}
ss << "}";
return ss.str();
}
TensorMap TensorMap::fromNpyFolder(const std::string& base_folder)
{
#if !defined(_WIN32)
DIR* dir_p = opendir(base_folder.c_str());
TLLM_CHECK_WITH_INFO(dir_p != nullptr, fmtstr("Could not open folder %s. ", base_folder.c_str()));
struct dirent* dp;
TensorMap ret_tensor;
while ((dp = readdir(dir_p)) != nullptr)
{
std::string filename(dp->d_name);
size_t len = filename.length();
if (len < 4 || filename.compare(len - 4, 4, ".npy"))
{
continue;
}
size_t pos = filename.find('-');
TLLM_CHECK_WITH_INFO(pos != std::string::npos, fmtstr("Invalid filename: %s\n", filename.c_str()));
MemoryType where;
if (filename.compare(0, pos, "GPU") == 0)
{
where = MEMORY_GPU;
}
else if (filename.compare(0, pos, "CPU") == 0)
{
where = MEMORY_CPU;
}
else if (filename.compare(0, pos, "CPU_PINNED") == 0)
{
where = MEMORY_CPU_PINNED;
}
else
{
TLLM_CHECK_WITH_INFO(false, fmtstr("Invalid filename: %s\n", filename.c_str()));
}
std::string key = filename.substr(pos + 1, len - pos - 5);
ret_tensor.tensor_map_.insert({key, Tensor::loadNpy(base_folder + "/" + filename, where)});
}
closedir(dir_p);
return ret_tensor;
#else
throw std::runtime_error("TensorMap::fromNpyFolder is not implemented on Windows.");
return {};
#endif // !defined(_WIN32)
}
void TensorMap::saveNpy(const std::string& base_folder)
{
#if !defined(_WIN32)
mode_t mode_0755 = S_IRWXU | S_IRGRP | S_IXGRP | S_IROTH | S_IXOTH;
int ret = mkdir(base_folder.c_str(), mode_0755);
TLLM_CHECK_WITH_INFO(ret == 0 || errno == EEXIST, fmtstr("Could not create folder %s.\n", base_folder.c_str()));
for (const auto& item : tensor_map_)
{
item.second.saveNpy(base_folder + "/" + item.second.whereToString() + "-" + item.first + ".npy");
}
#else
throw std::runtime_error("TensorMap::saveNpy is not implemented on Windows.");
#endif // !defined(_WIN32)
}
tensorConversion
inline DataType toTllmDataType(nvinfer1::DataType type)
{
switch (type)
{
case nvinfer1::DataType::kFLOAT: return DataType::TYPE_FP32;
case nvinfer1::DataType::kHALF: return DataType::TYPE_FP16;
case nvinfer1::DataType::kBF16: return DataType::TYPE_BF16;
case nvinfer1::DataType::kFP8: return DataType::TYPE_FP8_E4M3;
case nvinfer1::DataType::kINT8: return DataType::TYPE_INT8;
case nvinfer1::DataType::kUINT8: return DataType::TYPE_UINT8;
case nvinfer1::DataType::kINT32: return DataType::TYPE_INT32;
case nvinfer1::DataType::kINT64: return DataType::TYPE_INT64;
case nvinfer1::DataType::kBOOL: return DataType::TYPE_BOOL;
default: TLLM_THROW("Unsupported data type: %d", static_cast<int>(type));
}
}
inline MemoryType toTllmMemoryType(runtime::MemoryType type)
{
switch (type)
{
case runtime::MemoryType::kGPU: return MemoryType::MEMORY_GPU;
case runtime::MemoryType::kCPU: return MemoryType::MEMORY_CPU;
case runtime::MemoryType::kPINNED: return MemoryType::MEMORY_CPU_PINNED;
default: TLLM_THROW("Unsupported memory type: %d", static_cast<int>(type));
}
}
inline Tensor toTllmTensor(runtime::ITensor const& tensor)
{
MemoryType memoryType = toTllmMemoryType(tensor.getMemoryType());
DataType dataType = toTllmDataType(tensor.getDataType());
auto const& dims = tensor.getShape();
std::vector<std::size_t> shape(dims.d, dims.d + dims.nbDims);
auto* data = tensor.data();
return Tensor(memoryType, dataType, shape, data);
}
inline Tensor toTllmTensor(runtime::IBuffer const& buffer)
{
MemoryType memoryType = toTllmMemoryType(buffer.getMemoryType());
DataType dataType = toTllmDataType(buffer.getDataType());
std::vector<std::size_t> shape{buffer.getSize()};
auto* data = buffer.data();
return Tensor(memoryType, dataType, shape, data);
}
template <typename T>
Tensor toTllmTensor(MemoryType memoryType, std::vector<std::size_t> const& shape, T* data)
{
return Tensor{memoryType, getTensorType<T>(), shape, data};
}
template <typename T>
Tensor toTllmTensor(std::vector<T> const& data)
{
return Tensor{MemoryType::MEMORY_CPU, getTensorType<T>(), {data.size()}, data.data()};
}
template <typename T>
Tensor scalarToTllmTensor(T& data)
{
return Tensor{MemoryType::MEMORY_CPU, getTensorType<T>(), {1}, &data};
}
参考文献
- • https://github.com/NVIDIA/TensorRT-LLM/blob/release/0.5.0/cpp/tensorrt_llm/common/tensor.h
夜雨聆风
