chore: import upstream snapshot with attribution

samples/common/sampleDevice.h

@@ -0,0 +1,633 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef TRT_SAMPLE_DEVICE_H
+#define TRT_SAMPLE_DEVICE_H
+
+#include <cassert>
+#include <cstdint>
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <iostream>
+#include <thread>
+
+#include "common.h"
+#include "globalTimerKernel.h"
+#include "sampleUtils.h"
+
+namespace sample
+{
+
+//! True when Confidential Compute is enabled on the current system. Cached on
+//! the first call. When true, TrtCudaEvent falls back to a GPU global-timer
+//! kernel because cudaEventElapsedTime() is unreliable under CC (see nvbug
+//! 5598617, mirrors TRT-LLM PR #11657).
+[[nodiscard]] bool isConfidentialComputeEnabled();
+
+class TrtCudaEvent;
+
+namespace
+{
+
+void cudaSleep(void* sleep)
+{
+    std::this_thread::sleep_for(std::chrono::duration<float, std::milli>(*static_cast<float*>(sleep)));
+}
+
+} // namespace
+
+//!
+//! \class TrtCudaStream
+//! \brief Managed CUDA stream
+//!
+class TrtCudaStream
+{
+public:
+    TrtCudaStream()
+    {
+        CHECK(cudaStreamCreate(&mStream));
+    }
+
+    TrtCudaStream(const TrtCudaStream&) = delete;
+
+    TrtCudaStream& operator=(const TrtCudaStream&) = delete;
+
+    TrtCudaStream(TrtCudaStream&&) = delete;
+
+    TrtCudaStream& operator=(TrtCudaStream&&) = delete;
+
+    ~TrtCudaStream()
+    {
+        CHECK(cudaStreamDestroy(mStream));
+    }
+
+    cudaStream_t get() const
+    {
+        return mStream;
+    }
+
+    void synchronize()
+    {
+        CHECK(cudaStreamSynchronize(mStream));
+    }
+
+    void wait(TrtCudaEvent& event);
+
+    void sleep(float* ms)
+    {
+        CHECK(cudaLaunchHostFunc(mStream, cudaSleep, ms));
+    }
+
+private:
+    cudaStream_t mStream{};
+};
+
+//!
+//! \class TrtCudaEvent
+//! \brief Managed CUDA event
+//!
+class TrtCudaEvent
+{
+public:
+    explicit TrtCudaEvent(bool blocking = true)
+    {
+        const uint32_t flags = blocking ? cudaEventBlockingSync : cudaEventDefault;
+        CHECK(cudaEventCreateWithFlags(&mEvent, flags));
+        if (isConfidentialComputeEnabled())
+        {
+            CHECK(cudaMalloc(&mDeviceTimestamp, sizeof(uint64_t)));
+        }
+    }
+
+    TrtCudaEvent(TrtCudaEvent const&) = delete;
+
+    TrtCudaEvent& operator=(TrtCudaEvent const&) = delete;
+
+    TrtCudaEvent(TrtCudaEvent&&) = delete;
+
+    TrtCudaEvent& operator=(TrtCudaEvent&&) = delete;
+
+    ~TrtCudaEvent()
+    {
+        CHECK(cudaEventDestroy(mEvent));
+        if (mDeviceTimestamp != nullptr)
+        {
+            CHECK(cudaFree(mDeviceTimestamp));
+        }
+    }
+
+    cudaEvent_t get() const
+    {
+        return mEvent;
+    }
+
+    void record(TrtCudaStream const& stream)
+    {
+        if (mDeviceTimestamp != nullptr)
+        {
+            CHECK(launchGlobalTimerKernel(mDeviceTimestamp, stream.get()));
+        }
+        CHECK(cudaEventRecord(mEvent, stream.get()));
+    }
+
+    void synchronize() const
+    {
+        CHECK(cudaEventSynchronize(mEvent));
+    }
+
+    // Returns time elapsed time in milliseconds
+    float operator-(const TrtCudaEvent& e) const
+    {
+        // Synchronize both events to ensure they have completed before calculating elapsed time
+        synchronize();
+        e.synchronize();
+
+        if (mDeviceTimestamp != nullptr && e.mDeviceTimestamp != nullptr)
+        {
+            // Confidential Compute path: read %globaltimer values captured at record().
+            // cudaEventElapsedTime() is unreliable under CC (nvbug 5598617); the global
+            // timer kernel reads the same underlying register directly.
+            // Use signed int64_t so the subtraction is well-defined if the events
+            // are ever measured out of order (otherwise the unsigned->signed cast
+            // is implementation-defined in C++17).
+            int64_t endNs{0};
+            int64_t startNs{0};
+            CHECK(cudaMemcpy(&endNs, mDeviceTimestamp, sizeof(int64_t), cudaMemcpyDeviceToHost));
+            CHECK(cudaMemcpy(&startNs, e.mDeviceTimestamp, sizeof(int64_t), cudaMemcpyDeviceToHost));
+            return static_cast<float>(endNs - startNs) / 1.0e6F;
+        }
+
+        float time{0};
+        CHECK(cudaEventElapsedTime(&time, e.get(), get()));
+        return time;
+    }
+
+private:
+    cudaEvent_t mEvent{};
+    uint64_t* mDeviceTimestamp{nullptr};
+};
+
+inline void TrtCudaStream::wait(TrtCudaEvent& event)
+{
+    CHECK(cudaStreamWaitEvent(mStream, event.get(), 0));
+}
+
+//!
+//! \class TrtCudaGraph
+//! \brief Managed CUDA graph
+//!
+class TrtCudaGraph
+{
+public:
+    explicit TrtCudaGraph() = default;
+
+    TrtCudaGraph(const TrtCudaGraph&) = delete;
+
+    TrtCudaGraph& operator=(const TrtCudaGraph&) = delete;
+
+    TrtCudaGraph(TrtCudaGraph&&) = delete;
+
+    TrtCudaGraph& operator=(TrtCudaGraph&&) = delete;
+
+    ~TrtCudaGraph()
+    {
+        if (mGraphExec)
+        {
+            cudaGraphExecDestroy(mGraphExec);
+        }
+    }
+
+    void beginCapture(TrtCudaStream& stream)
+    {
+        CHECK(cudaStreamBeginCapture(stream.get(), cudaStreamCaptureModeThreadLocal));
+    }
+
+    bool launch(TrtCudaStream& stream)
+    {
+        return cudaGraphLaunch(mGraphExec, stream.get()) == cudaSuccess;
+    }
+
+    void endCapture(TrtCudaStream& stream)
+    {
+        CHECK(cudaStreamEndCapture(stream.get(), &mGraph));
+        CHECK(cudaGraphInstantiate(&mGraphExec, mGraph, nullptr, nullptr, 0));
+        CHECK(cudaGraphDestroy(mGraph));
+    }
+
+    void endCaptureOnError(TrtCudaStream& stream)
+    {
+        // There are two possibilities why stream capture would fail:
+        // (1) stream is in cudaErrorStreamCaptureInvalidated state.
+        // (2) TRT reports a failure.
+        // In case (1), the returning mGraph should be nullptr.
+        // In case (2), the returning mGraph is not nullptr, but it should not be used.
+        const auto ret = cudaStreamEndCapture(stream.get(), &mGraph);
+        if (ret == cudaErrorStreamCaptureInvalidated)
+        {
+            assert(mGraph == nullptr);
+        }
+        else
+        {
+            CHECK(ret);
+            assert(mGraph != nullptr);
+            CHECK(cudaGraphDestroy(mGraph));
+            mGraph = nullptr;
+        }
+        // Clean up any CUDA error.
+        cudaGetLastError();
+        sample::gLogWarning << "The CUDA graph capture on the stream has failed." << std::endl;
+    }
+
+private:
+    cudaGraph_t mGraph{};
+    cudaGraphExec_t mGraphExec{};
+};
+
+//!
+//! \class TrtCudaBuffer
+//! \brief Managed buffer for host and device
+//!
+template <typename A, typename D>
+class TrtCudaBuffer
+{
+public:
+    TrtCudaBuffer() = default;
+
+    TrtCudaBuffer(const TrtCudaBuffer&) = delete;
+
+    TrtCudaBuffer& operator=(const TrtCudaBuffer&) = delete;
+
+    TrtCudaBuffer(TrtCudaBuffer&& rhs)
+    {
+        reset(rhs.mPtr, rhs.mSize);
+        rhs.mPtr = nullptr;
+        rhs.mSize = 0;
+    }
+
+    TrtCudaBuffer& operator=(TrtCudaBuffer&& rhs)
+    {
+        if (this != &rhs)
+        {
+            reset(rhs.mPtr, rhs.mSize);
+            rhs.mPtr = nullptr;
+            rhs.mSize = 0;
+        }
+        return *this;
+    }
+
+    ~TrtCudaBuffer()
+    {
+        reset();
+    }
+
+    TrtCudaBuffer(size_t size)
+    {
+        A()(&mPtr, size);
+        mSize = size;
+    }
+
+    void allocate(size_t size)
+    {
+        reset();
+        A()(&mPtr, size);
+        mSize = size;
+    }
+
+    void reset(void* ptr = nullptr, size_t size = 0)
+    {
+        if (mPtr)
+        {
+            D()(mPtr);
+        }
+        mPtr = ptr;
+        mSize = size;
+    }
+
+    void* get() const
+    {
+        return mPtr;
+    }
+
+    size_t getSize() const
+    {
+        return mSize;
+    }
+
+private:
+    void* mPtr{nullptr};
+    size_t mSize{0};
+};
+
+struct DeviceAllocator
+{
+    void operator()(void** ptr, size_t size)
+    {
+        CHECK(cudaMalloc(ptr, size));
+    }
+};
+
+struct DeviceDeallocator
+{
+    void operator()(void* ptr)
+    {
+        CHECK(cudaFree(ptr));
+    }
+};
+
+struct ManagedAllocator
+{
+    void operator()(void** ptr, size_t size)
+    {
+        CHECK(cudaMallocManaged(ptr, size));
+    }
+};
+
+struct HostAllocator
+{
+    //! Attempts to allocate size bytes on host, pointing *ptr to the start.
+    //! First attempts to allocate pinned memory using cudaMallocHost(ptr, size), failing that, warns to gLogWarning and
+    //! falls back to ::operator new(size) to allocate pageable memory. If that still fails, an exception may be thrown.
+    void operator()(void** ptr, size_t size)
+    {
+        // Try allocating pinned host memory.
+        cudaError_t ret = cudaMallocHost(ptr, size);
+
+        // If we cannot allocate pinned host memory, allocate pageable host memory instead and print a warning.
+        if (ret != cudaSuccess)
+        {
+            // Clean up the last cuda error.
+            (void) cudaGetLastError();
+
+            sample::gLogWarning << "cudaMallocHost() call with ptr=" << ptr << " and size=" << size
+                                << " returns a cuda error: " << cudaGetErrorString(ret) << std::endl;
+            sample::gLogWarning << "Allocate pageable host memory instead of pinned host memory. H2D and D2H copy "
+                                   "latencies may become longer."
+                                << std::endl;
+            *ptr = ::operator new(size);
+
+            // Make sure there is no remaining cuda error at this point.
+            CHECK(cudaGetLastError());
+        }
+    }
+};
+
+struct HostDeallocator
+{
+    //! Attempts to deallocate the host memory allocated by HostAllocator.
+    //! It first checks if ptr is a pinned or pageable host memory. If pinned, call cudaFreeHost() to free it. If
+    //! pageable, call ::operator delete() to free it. If ptr is neither of them, an error is printed and the program
+    //! exits.
+    void operator()(void* ptr)
+    {
+        // Check if the host memory pointer is pinned or pageable.
+        cudaPointerAttributes attrs;
+        CHECK(cudaPointerGetAttributes(&attrs, ptr));
+
+        // If pinned, call cudaFreeHost() to deallocate it. Under Confidential
+        // Compute, memory returned by cudaMallocHost() may be reported as
+        // cudaMemoryTypeManaged; it must still be released via cudaFreeHost.
+        if (attrs.type == cudaMemoryTypeHost || attrs.type == cudaMemoryTypeManaged)
+        {
+            CHECK(cudaFreeHost(ptr));
+        }
+        // If pageable, delete it directly.
+        else if (attrs.type == cudaMemoryTypeUnregistered)
+        {
+            ::operator delete(ptr);
+        }
+        // The host memory pointer should not be of any other types.
+        else
+        {
+            sample::gLogError << "Unexpected cuda memory type:" << static_cast<int32_t>(attrs.type) << std::endl;
+            exit(EXIT_FAILURE);
+        }
+    }
+};
+
+using TrtDeviceBuffer = TrtCudaBuffer<DeviceAllocator, DeviceDeallocator>;
+using TrtManagedBuffer = TrtCudaBuffer<ManagedAllocator, DeviceDeallocator>;
+
+using TrtHostBuffer = TrtCudaBuffer<HostAllocator, HostDeallocator>;
+
+//!
+//! \class MirroredBuffer
+//! \brief Coupled host and device buffers
+//!
+class IMirroredBuffer
+{
+public:
+    //!
+    //! Allocate memory for the mirrored buffer give the size
+    //! of the allocation.
+    //!
+    virtual void allocate(size_t size) = 0;
+
+    //!
+    //! Get the pointer to the device side buffer.
+    //!
+    //! \return pointer to device memory or nullptr if uninitialized.
+    //!
+    virtual void* getDeviceBuffer() const = 0;
+
+    //!
+    //! Get the pointer to the host side buffer.
+    //!
+    //! \return pointer to host memory or nullptr if uninitialized.
+    //!
+    virtual void* getHostBuffer() const = 0;
+
+    //!
+    //! Copy the memory from host to device.
+    //!
+    virtual void hostToDevice(TrtCudaStream& stream) = 0;
+
+    //!
+    //! Copy the memory from device to host.
+    //!
+    virtual void deviceToHost(TrtCudaStream& stream) = 0;
+
+    //!
+    //! Interface to get the size of the memory
+    //!
+    //! \return the size of memory allocated.
+    //!
+    virtual size_t getSize() const = 0;
+
+    //!
+    //! Virtual destructor declaraion
+    //!
+    virtual ~IMirroredBuffer() = default;
+
+}; // class IMirroredBuffer
+
+//!
+//! Class to have a separate memory buffer for discrete device and host allocations.
+//!
+class DiscreteMirroredBuffer : public IMirroredBuffer
+{
+public:
+    void allocate(size_t size) override
+    {
+        mSize = size;
+        mHostBuffer.allocate(size);
+        mDeviceBuffer.allocate(size);
+    }
+
+    void* getDeviceBuffer() const override
+    {
+        return mDeviceBuffer.get();
+    }
+
+    void* getHostBuffer() const override
+    {
+        return mHostBuffer.get();
+    }
+
+    void hostToDevice(TrtCudaStream& stream) override
+    {
+        CHECK(cudaMemcpyAsync(mDeviceBuffer.get(), mHostBuffer.get(), mSize, cudaMemcpyHostToDevice, stream.get()));
+    }
+
+    void deviceToHost(TrtCudaStream& stream) override
+    {
+        CHECK(cudaMemcpyAsync(mHostBuffer.get(), mDeviceBuffer.get(), mSize, cudaMemcpyDeviceToHost, stream.get()));
+    }
+
+    size_t getSize() const override
+    {
+        return mSize;
+    }
+
+private:
+    size_t mSize{0};
+    TrtHostBuffer mHostBuffer;
+    TrtDeviceBuffer mDeviceBuffer;
+}; // class DiscreteMirroredBuffer
+
+//!
+//! Class to have a unified memory buffer for embedded devices.
+//!
+class UnifiedMirroredBuffer : public IMirroredBuffer
+{
+public:
+    void allocate(size_t size) override
+    {
+        mSize = size;
+        mBuffer.allocate(size);
+    }
+
+    void* getDeviceBuffer() const override
+    {
+        return mBuffer.get();
+    }
+
+    void* getHostBuffer() const override
+    {
+        return mBuffer.get();
+    }
+
+    void hostToDevice(TrtCudaStream& stream) override
+    {
+        // Does nothing since we are using unified memory.
+    }
+
+    void deviceToHost(TrtCudaStream& stream) override
+    {
+        // Does nothing since we are using unified memory.
+    }
+
+    size_t getSize() const override
+    {
+        return mSize;
+    }
+
+private:
+    size_t mSize{0};
+    TrtManagedBuffer mBuffer;
+}; // class UnifiedMirroredBuffer
+
+//!
+//! Class to allocate memory for outputs with data-dependent shapes. The sizes of those are unknown so pre-allocation is
+//! not possible.
+//!
+class OutputAllocator : public nvinfer1::IOutputAllocator
+{
+public:
+    //! Construct, using buffer as the backing storage:
+    explicit OutputAllocator(std::shared_ptr<IMirroredBuffer> buffer)
+        : mBuffer{std::move(buffer)}
+    {
+        ASSERT(mBuffer);
+    }
+
+    ~OutputAllocator() override = default;
+
+    void* reallocateOutput(
+        char const* tensorName, void* currentMemory, uint64_t size, uint64_t alignment) noexcept override
+    {
+        // Some memory allocators return nullptr when allocating zero bytes, but TensorRT requires a non-null ptr
+        // even for empty tensors, so allocate a dummy byte.
+        size = std::max(size, static_cast<uint64_t>(1));
+        if (size > mSize)
+        {
+            mBuffer->allocate(roundUp(size, alignment));
+            mSize = size;
+        }
+        return mBuffer->getDeviceBuffer();
+    }
+
+    //! IMirroredBuffer does not implement Async allocation, hence this is just a wrap around
+    void* reallocateOutputAsync(char const* tensorName, void* currentMemory, uint64_t size, uint64_t alignment,
+        cudaStream_t /*stream*/) noexcept override
+    {
+        return reallocateOutput(tensorName, currentMemory, size, alignment);
+    }
+
+    void notifyShape(char const* tensorName, nvinfer1::Dims const& dims) noexcept override
+    {
+        mFinalDims = dims;
+    }
+
+    IMirroredBuffer* getBuffer()
+    {
+        return mBuffer.get();
+    }
+
+    nvinfer1::Dims getFinalDims()
+    {
+        return mFinalDims;
+    }
+
+private:
+    std::shared_ptr<IMirroredBuffer> mBuffer;
+    uint64_t mSize{};
+    nvinfer1::Dims mFinalDims;
+};
+
+//! Set the GPU to run the inference on.
+void setCudaDevice(int32_t device, std::ostream& os);
+
+//! Get the CUDA version of the current CUDA driver.
+int32_t getCudaDriverVersion();
+
+//! Get the CUDA version of the current CUDA runtime.
+int32_t getCudaRuntimeVersion();
+
+
+} // namespace sample
+
+#endif // TRT_SAMPLE_DEVICE_H