diff --git a/onnxruntime/core/providers/cuda/cuda_utils.cu b/onnxruntime/core/providers/cuda/cuda_utils.cu index 01bb6a73df..d6cdebb17c 100644 --- a/onnxruntime/core/providers/cuda/cuda_utils.cu +++ b/onnxruntime/core/providers/cuda/cuda_utils.cu @@ -3,10 +3,6 @@ // Thrust code needs to be compiled with nvcc #include -#include -#include -#include -#include #include "core/providers/cuda/shared_inc/cuda_utils.h" #include "core/providers/cuda/cu_inc/common.cuh" #include "cudnn_common.h" @@ -14,19 +10,44 @@ namespace onnxruntime { namespace cuda { +template +__global__ void _Fill( + T* output_data, + T val, + CUDA_LONG N) { + CALCULATE_ELEMENTWISE_INDEX_OR_EXIT(id, N); + output_data[id] = val; +} + template class ConstantBufferImpl : public IConstantBuffer { public: - ConstantBufferImpl(T val) : val_(val) {} + ConstantBufferImpl(T val) : val_(val), buffer_(nullptr), count_(0) { + } + ~ConstantBufferImpl() { + if (buffer_) + cudaFree(buffer_); + } virtual const T* GetBuffer(size_t count) { - buffer_.resize(count); - thrust::fill(buffer_.begin(), buffer_.end(), val_); - return buffer_.data().get(); + if (count > count_) { + if (buffer_) { + cudaFree(buffer_); + buffer_ = nullptr; + } + CUDA_CALL_THROW(cudaMalloc(&buffer_, count * sizeof(T))); + count_ = count; + + int blocksPerGrid = (int)(ceil(static_cast(count) / GridDim::maxThreadsPerBlock)); + CUDA_LONG N = static_cast(count); + _Fill<<>>(buffer_, val_, N); + } + return buffer_; } private: - thrust::device_vector buffer_; + T* buffer_; + size_t count_; T val_; };