Skip to content

fix: Support stream readers before TensorRT 10.7 #147

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
suharvest wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix: Support stream readers before TensorRT 10.7 #147

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 29 additions & 8 deletions cpp/common/trtUtils.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@ namespace trt_edgellm
namespace
{

#if NV_TENSORRT_MAJOR >= 11 || (NV_TENSORRT_MAJOR == 10 && NV_TENSORRT_MINOR >= 7)
constexpr std::size_t kEngineDeviceTransferChunkSize = 4U * 1024U * 1024U;
constexpr std::size_t kEngineDeviceTransferBufferCount = 2U;

Expand Down Expand Up @@ -151,14 +152,22 @@ class PinnedHostBuffer
cudaEvent_t mEvent{nullptr};
bool mCopyPending{false};
};
#endif

class FileStreamReaderV2 : public nvinfer1::IStreamReaderV2
#if NV_TENSORRT_MAJOR >= 11 || (NV_TENSORRT_MAJOR == 10 && NV_TENSORRT_MINOR >= 7)
class FileStreamReader : public nvinfer1::IStreamReaderV2
#else
class FileStreamReader : public nvinfer1::IStreamReader
#endif
{
public:
explicit FileStreamReaderV2(std::filesystem::path const& filePath)
explicit FileStreamReader(std::filesystem::path const& filePath)
: mPath(filePath.string())
, mDeviceTransferBuffers{
PinnedHostBuffer{kEngineDeviceTransferChunkSize}, PinnedHostBuffer{kEngineDeviceTransferChunkSize}}
#if NV_TENSORRT_MAJOR >= 11 || (NV_TENSORRT_MAJOR == 10 && NV_TENSORRT_MINOR >= 7)
, mDeviceTransferBuffers{PinnedHostBuffer{kEngineDeviceTransferChunkSize}, PinnedHostBuffer {
kEngineDeviceTransferChunkSize
}}
#endif
{
mFd = open(mPath.c_str(), O_RDONLY);
if (mFd < 0)
Expand All @@ -182,18 +191,22 @@ class FileStreamReaderV2 : public nvinfer1::IStreamReaderV2
mSize = static_cast<int64_t>(status.st_size);
}

FileStreamReaderV2(FileStreamReaderV2 const&) = delete;
FileStreamReaderV2& operator=(FileStreamReaderV2 const&) = delete;
FileStreamReader(FileStreamReader const&) = delete;
FileStreamReader& operator=(FileStreamReader const&) = delete;

~FileStreamReaderV2() override
~FileStreamReader() override
{
if (mFd >= 0)
{
close(mFd);
}
}

#if NV_TENSORRT_MAJOR >= 11 || (NV_TENSORRT_MAJOR == 10 && NV_TENSORRT_MINOR >= 7)
int64_t read(void* destination, int64_t nbBytes, cudaStream_t stream) noexcept override
#else
int64_t read(void* destination, int64_t nbBytes) override
#endif
{
if (destination == nullptr || nbBytes < 0)
{
Expand All @@ -205,13 +218,16 @@ class FileStreamReaderV2 : public nvinfer1::IStreamReaderV2
}

int64_t const bytesToRead = std::min(nbBytes, mSize - mOffset);
#if NV_TENSORRT_MAJOR >= 11 || (NV_TENSORRT_MAJOR == 10 && NV_TENSORRT_MINOR >= 7)
if (isDeviceAccessiblePointer(destination))
{
return readToDevice(destination, bytesToRead, stream);
}
#endif
return readToHost(destination, bytesToRead);
}

#if NV_TENSORRT_MAJOR >= 11 || (NV_TENSORRT_MAJOR == 10 && NV_TENSORRT_MINOR >= 7)
bool seek(int64_t offset, nvinfer1::SeekPosition where) noexcept override
{
int64_t base = 0;
Expand All @@ -231,6 +247,7 @@ class FileStreamReaderV2 : public nvinfer1::IStreamReaderV2
mOffset = nextOffset;
return true;
}
#endif

private:
int64_t readToHost(void* destination, int64_t bytesToRead) noexcept
Expand Down Expand Up @@ -259,6 +276,7 @@ class FileStreamReaderV2 : public nvinfer1::IStreamReaderV2
return totalRead;
}

#if NV_TENSORRT_MAJOR >= 11 || (NV_TENSORRT_MAJOR == 10 && NV_TENSORRT_MINOR >= 7)
int64_t readToDevice(void* destination, int64_t bytesToRead, cudaStream_t stream) noexcept
{
auto* output = static_cast<std::byte*>(destination);
Expand Down Expand Up @@ -296,12 +314,15 @@ class FileStreamReaderV2 : public nvinfer1::IStreamReaderV2
}
return totalRead;
}
#endif

std::string mPath;
int mFd{-1};
int64_t mSize{0};
int64_t mOffset{0};
#if NV_TENSORRT_MAJOR >= 11 || (NV_TENSORRT_MAJOR == 10 && NV_TENSORRT_MINOR >= 7)
std::array<PinnedHostBuffer, kEngineDeviceTransferBufferCount> mDeviceTransferBuffers;
#endif
};

} // namespace
Expand Down Expand Up @@ -461,7 +482,7 @@ bool engineHasOutputTensor(nvinfer1::ICudaEngine const* engine, char const* tens
std::unique_ptr<nvinfer1::ICudaEngine> deserializeCudaEngineFromFile(
nvinfer1::IRuntime& runtime, std::filesystem::path const& enginePath)
{
FileStreamReaderV2 streamReader(enginePath);
FileStreamReader streamReader(enginePath);
auto engine = std::unique_ptr<nvinfer1::ICudaEngine>(runtime.deserializeCudaEngine(streamReader));
if (!engine)
{
Expand Down