wehub-resource-sync
137 lines
4.5 KiB
C++
137 lines
4.5 KiB
C++
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
|
|
//
|
|
// 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.
|
|
|
|
// The file has been adapted from pytorch project
|
|
// Licensed under BSD-style license -
|
|
// https://github.com/pytorch/pytorch/blob/main/LICENSE
|
|
|
|
#pragma once
|
|
// Light-weight version of CUDAContext.h with fewer transitive includes
|
|
|
|
// cublasLT was introduced in CUDA 10.1 but we enable only for 11.1 that also
|
|
// added bf16 support
|
|
|
|
#if defined(PADDLE_WITH_HIP)
|
|
#include <hip/hip_runtime.h>
|
|
#elif defined(PADDLE_WITH_CUDA)
|
|
#if defined(USE_CUDSS)
|
|
#include <cudss.h>
|
|
#endif
|
|
#include <driver_types.h>
|
|
#endif
|
|
|
|
#include <c10/core/Allocator.h>
|
|
#include <c10/cuda/CUDAFunctions.h>
|
|
|
|
#include <cstdint>
|
|
#include <map>
|
|
#include <shared_mutex>
|
|
#include <tuple>
|
|
|
|
#include "paddle/common/macros.h"
|
|
#include "paddle/phi/backends/gpu/forwards.h"
|
|
|
|
namespace c10 {
|
|
struct Allocator;
|
|
}
|
|
|
|
namespace at::cuda {
|
|
|
|
#if defined(PADDLE_WITH_HIP)
|
|
using CUDAContextDeviceProp = phi::gpuDeviceProp;
|
|
using CUDAContextSparseHandle = phi::sparseHandle_t;
|
|
using CUDAContextBlasHandle = phi::blasHandle_t;
|
|
using CUDAContextBlasLtHandle = phi::blasLtHandle_t;
|
|
using CUDAContextSolverHandle = phi::solverHandle_t;
|
|
#elif defined(PADDLE_WITH_CUDA)
|
|
using CUDAContextDeviceProp = cudaDeviceProp;
|
|
using CUDAContextSparseHandle = cusparseHandle_t;
|
|
using CUDAContextBlasHandle = cublasHandle_t;
|
|
using CUDAContextBlasLtHandle = cublasLtHandle_t;
|
|
using CUDAContextSolverHandle = cusolverDnHandle_t;
|
|
#endif
|
|
|
|
/*
|
|
A common CUDA interface for ATen.
|
|
|
|
This interface is distinct from CUDAHooks, which defines an interface that links
|
|
to both CPU-only and CUDA builds. That interface is intended for runtime
|
|
dispatch and should be used from files that are included in both CPU-only and
|
|
CUDA builds.
|
|
|
|
CUDAContext, on the other hand, should be preferred by files only included in
|
|
CUDA builds. It is intended to expose CUDA functionality in a consistent
|
|
manner.
|
|
|
|
This means there is some overlap between the CUDAContext and CUDAHooks, but
|
|
the choice of which to use is simple: use CUDAContext when in a CUDA-only file,
|
|
use CUDAHooks otherwise.
|
|
|
|
Note that CUDAContext simply defines an interface with no associated class.
|
|
It is expected that the modules whose functions compose this interface will
|
|
manage their own state. There is only a single CUDA context/state.
|
|
*/
|
|
|
|
/**
|
|
* DEPRECATED: use device_count() instead
|
|
*/
|
|
inline int64_t getNumGPUs() { return c10::cuda::device_count(); }
|
|
|
|
/**
|
|
* CUDA is available if we compiled with CUDA, and there are one or more
|
|
* devices. If we compiled with CUDA but there is a driver problem, etc.,
|
|
* this function will report CUDA is not available (rather than raise an error.)
|
|
*/
|
|
inline bool is_available() { return c10::cuda::device_count() > 0; }
|
|
|
|
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
|
PADDLE_API CUDAContextDeviceProp* getCurrentDeviceProperties();
|
|
|
|
PADDLE_API int warp_size();
|
|
|
|
PADDLE_API CUDAContextDeviceProp* getDeviceProperties(c10::DeviceIndex device);
|
|
|
|
PADDLE_API bool canDeviceAccessPeer(c10::DeviceIndex device,
|
|
c10::DeviceIndex peer_device);
|
|
|
|
/* Handles */
|
|
PADDLE_API CUDAContextSparseHandle getCurrentCUDASparseHandle();
|
|
PADDLE_API CUDAContextBlasHandle getCurrentCUDABlasHandle();
|
|
PADDLE_API CUDAContextBlasLtHandle getCurrentCUDABlasLtHandle();
|
|
|
|
PADDLE_API void clearCublasWorkspaces();
|
|
struct WorkspaceMapWithMutex {
|
|
std::map<std::tuple<void*, void*>, at::DataPtr> map;
|
|
std::shared_mutex mutex;
|
|
};
|
|
|
|
PADDLE_API WorkspaceMapWithMutex& cublas_handle_stream_to_workspace();
|
|
PADDLE_API WorkspaceMapWithMutex& cublaslt_handle_stream_to_workspace();
|
|
PADDLE_API size_t getChosenWorkspaceSize();
|
|
PADDLE_API size_t getCUDABlasLtWorkspaceSize();
|
|
PADDLE_API void* getCUDABlasLtWorkspace();
|
|
|
|
PADDLE_API CUDAContextSolverHandle getCurrentCUDASolverDnHandle();
|
|
|
|
#if defined(USE_CUDSS)
|
|
PADDLE_API cudssHandle_t getCurrentCudssHandle();
|
|
#endif
|
|
|
|
// Get the CUDA device allocator for the current device.
|
|
// Returns a pointer to a c10::Allocator that allocates GPU memory.
|
|
PADDLE_API c10::Allocator* getCUDADeviceAllocator();
|
|
#endif
|
|
|
|
} // namespace at::cuda
|