sources/includes/ds3d/common/helper/cuda_utils.h

Go to the documentation of this file.

/*
 * SPDX-FileCopyrightText: Copyright (c) 2023-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 DS3D_COMMON_HELPER_CUDA_UTILS_H
#define DS3D_COMMON_HELPER_CUDA_UTILS_H
 
#include "ds3d/common/common.h"
#include "ds3d/common/helper/memdata.h"
 
#include <cuda.h>
#include <cuda_runtime_api.h>
 
namespace ds3d {
 
class CudaStream {
    cudaStream_t _stream = nullptr;
    int _gpuId = 0;
    DS3D_DISABLE_CLASS_COPY(CudaStream);
 
public:
    explicit CudaStream(uint flag = cudaStreamDefault, int gpuId = 0, int priority = 0)
    {
        cudaSetDevice(_gpuId);
        DS3D_CHECK_CUDA_ERROR(
            cudaStreamCreateWithPriority(&_stream, flag, priority), , "cudaStreamCreate failed");
    }
    ~CudaStream()
    {
        if (_stream != nullptr) {
            DS3D_CHECK_CUDA_ERROR(
                cudaSetDevice(_gpuId), , "cudaStreamDestroy failed to set gpu-id:%d", _gpuId);
            DS3D_CHECK_CUDA_ERROR(cudaStreamDestroy(_stream), , "cudaStreamDestroy failed");
        }
    }
    ErrCode sync()
    {
        if (!_stream) {
            return ErrCode::kGood;
        }
        DS3D_CHECK_CUDA_ERROR(
            cudaSetDevice(_gpuId), return ErrCode::kCuda,
            "cudaStreamSynchronize failed to set gpu-id:%d", _gpuId);
        DS3D_CHECK_CUDA_ERROR(
            cudaStreamSynchronize(_stream), return ErrCode::kCuda, "cudaStreamSynchronize failed");
        return ErrCode::kGood;
    }
    int gpuId() const { return _gpuId; }
    cudaStream_t& get() { return _stream; }
};
 
class GpuCudaMemBuf : public MemData {
public:
    GpuCudaMemBuf() {}
    static std::unique_ptr<GpuCudaMemBuf> CreateBuf(size_t size, int gpuId)
    {
        auto mem = std::make_unique<GpuCudaMemBuf>();
 
        DS3D_CHECK_CUDA_ERROR(
            cudaSetDevice(gpuId), return nullptr,
            "cudaSetDevice: %d failed when allocate cuda memory", gpuId);
        size = DS3D_ROUND_UP(size, 256);
        void* data = nullptr;
        DS3D_CHECK_CUDA_ERROR(
            cudaMalloc(&data, size), return nullptr, "cudaMalloc size: %d failed", (int)size);
        DS_ASSERT(data);
 
        mem->data = data;
        mem->byteSize = size;
        mem->devId = 0;
        mem->type = MemType::kGpuCuda;
        return mem;
    }
    ~GpuCudaMemBuf()
    {
        if (data) {
            cudaSetDevice(devId);
            cudaFree(data);
            data = nullptr;
            byteSize = 0;
        }
    }
};
 
};  // namespace ds3d
#endif