wehub-resource-sync
115 lines
4.3 KiB
C++
115 lines
4.3 KiB
C++
// Copyright (c) 2024 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.
|
|
|
|
#pragma once
|
|
#ifdef PADDLE_WITH_CUDA
|
|
#include <cuda_runtime.h>
|
|
#endif
|
|
|
|
#include "paddle/cinn/adt/adt.h"
|
|
#include "paddle/cinn/common/macros.h"
|
|
#include "paddle/cinn/common/target.h"
|
|
#include "paddle/cinn/runtime/backend_api.h"
|
|
#include "paddle/common/enforce.h"
|
|
#include "paddle/phi/backends/device_manager.h"
|
|
#include "paddle/phi/common/place.h"
|
|
|
|
namespace cinn::runtime {
|
|
|
|
inline std::optional<int> GetArchDevice(const common::Target& target) {
|
|
return target.arch.Match(
|
|
[&](common::UnknownArch) -> std::optional<int> { return std::nullopt; },
|
|
[&](common::X86Arch) -> std::optional<int> { return std::nullopt; },
|
|
[&](common::ARMArch) -> std::optional<int> { return std::nullopt; },
|
|
[&](common::NVGPUArch) -> std::optional<int> {
|
|
#ifdef CINN_WITH_CUDA
|
|
int device_id;
|
|
PADDLE_ENFORCE_EQ(
|
|
cudaGetDevice(&device_id),
|
|
cudaSuccess,
|
|
::common::errors::InvalidArgument("cudaGetDevice failed!"));
|
|
return std::optional<int>{device_id};
|
|
#else
|
|
return std::nullopt;
|
|
#endif
|
|
},
|
|
[&](common::CustomDeviceArch) -> std::optional<int> {
|
|
#ifdef CINN_WITH_CUSTOM_DEVICE
|
|
int device_id = phi::DeviceManager::GetDevice(target.device_name_str());
|
|
return std::optional<int>{device_id};
|
|
#else
|
|
return std::nullopt;
|
|
#endif
|
|
},
|
|
[&](common::HygonDCUArchHIP) -> std::optional<int> {
|
|
int device_id =
|
|
BackendAPI::get_backend(common::HygonDCUArchHIP{})->get_device();
|
|
return std::optional<int>{device_id};
|
|
},
|
|
[&](common::HygonDCUArchSYCL) -> std::optional<int> {
|
|
int device_id =
|
|
BackendAPI::get_backend(common::HygonDCUArchSYCL{})->get_device();
|
|
return std::optional<int>{device_id};
|
|
});
|
|
}
|
|
|
|
inline void SetArchDevice(const common::Target& target,
|
|
const std::optional<int>& device_id) {
|
|
target.arch.Match(
|
|
[&](common::UnknownArch) -> void {},
|
|
[&](common::X86Arch) -> void {},
|
|
[&](common::ARMArch) -> void {},
|
|
[&](common::NVGPUArch) -> void {
|
|
#ifdef CINN_WITH_CUDA
|
|
PADDLE_ENFORCE_EQ(device_id.has_value(),
|
|
true,
|
|
::common::errors::InvalidArgument(
|
|
"Required device_id should have value, but "
|
|
"received std::nullopt."));
|
|
cudaSetDevice(device_id.value());
|
|
#endif
|
|
},
|
|
[&](common::CustomDeviceArch) -> void {
|
|
#ifdef CINN_WITH_CUSTOM_DEVICE
|
|
PADDLE_ENFORCE_EQ(device_id.has_value(),
|
|
true,
|
|
::common::errors::InvalidArgument(
|
|
"Required device_id should have value, but "
|
|
"received std::nullopt."));
|
|
phi::DeviceManager::SetDevice(target.device_name_str(),
|
|
device_id.value());
|
|
#endif
|
|
},
|
|
[&](common::HygonDCUArchHIP) -> void {
|
|
PADDLE_ENFORCE_EQ(device_id.has_value(),
|
|
true,
|
|
::common::errors::InvalidArgument(
|
|
"Required device_id should have value, but "
|
|
"received std::nullopt."));
|
|
BackendAPI::get_backend(common::HygonDCUArchHIP{})
|
|
->set_device(device_id.value());
|
|
},
|
|
[&](common::HygonDCUArchSYCL) -> void {
|
|
PADDLE_ENFORCE_EQ(device_id.has_value(),
|
|
true,
|
|
::common::errors::InvalidArgument(
|
|
"Required device_id should have value, but "
|
|
"received std::nullopt."));
|
|
BackendAPI::get_backend(common::HygonDCUArchSYCL{})
|
|
->set_device(device_id.value());
|
|
});
|
|
}
|
|
|
|
} // namespace cinn::runtime
|