mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 00:21:07 +01:00
b4173b80b7
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/73445 Refactors the whole codebase to use LazyTensorPtr (defined as c10::intrusive_ptr) to enable XLA to use a derived class XlaLazyTensor and override functionality. this PR is just the first step, and we will need to add a factory class that XLA can override in their backend to actually hook up their derived tensor class. Parallel PR on lazy_tensor_staging: #73429 Test Plan: tested via lazy_tensor_staging test_ptltc and torchbench and CI Reviewed By: ezyang Differential Revision: D34481918 fbshipit-source-id: 01176b127df6b79039aa1bc57bc6da5505161f87 (cherry picked from commit 52b9ae4e22d2703d44c6436311d79d40bd62c6aa)
70 lines
2.1 KiB
C++
70 lines
2.1 KiB
C++
#include <torch/csrc/lazy/backend/backend_device.h>
|
|
|
|
#include <c10/core/Device.h>
|
|
#include <c10/util/Exception.h>
|
|
#include <c10/util/StringUtil.h>
|
|
#include <torch/csrc/lazy/core/tensor.h>
|
|
#include <torch/csrc/lazy/backend/backend_interface.h>
|
|
|
|
namespace torch {
|
|
namespace lazy {
|
|
|
|
// TODO(alanwaketan): Use the backend API to get the default device type.
|
|
// In the future, we should also get the default device ordinal.
|
|
BackendDevice::BackendDevice()
|
|
: type_(std::make_shared<BackendDeviceType>()) {}
|
|
|
|
BackendDevice::BackendDevice(std::shared_ptr<BackendDeviceType>&& type, int64_t ordinal)
|
|
: type_(std::move(type)), ordinal_(ordinal) {}
|
|
|
|
BackendDevice::BackendDevice(const std::string& device_spec)
|
|
: BackendDevice::BackendDevice() {}
|
|
|
|
int8_t BackendDevice::type() const {
|
|
TORCH_INTERNAL_ASSERT(type_);
|
|
return type_->type;
|
|
}
|
|
|
|
std::string BackendDevice::toString() const {
|
|
TORCH_INTERNAL_ASSERT(type_);
|
|
return c10::str(type_->toString(), ordinal_);
|
|
}
|
|
|
|
int BackendDevice::compare(const BackendDevice& rhs) const {
|
|
if (type() != rhs.type()) {
|
|
return type() < rhs.type() ? -1 : +1;
|
|
}
|
|
return ordinal_ < rhs.ordinal_ ? -1 : (ordinal_ > rhs.ordinal_ ? +1 : 0);
|
|
}
|
|
|
|
std::ostream& operator<<(std::ostream& os, const BackendDevice& device) {
|
|
os << device.toString();
|
|
return os;
|
|
}
|
|
|
|
// TODO(whc) refactor this: we need to support non-zero default ordinal for torch/XLA.
|
|
BackendDevice atenDeviceToBackendDevice(const c10::Device& device) {
|
|
TORCH_CHECK(device.type() == at::kLazy, device);
|
|
int64_t ordinal = device.has_index() ? device.index() : 0;
|
|
return BackendDevice(getBackend()->GetDefaultDeviceType(), ordinal);
|
|
}
|
|
|
|
// TODO(whc) refactor this: we need to support non 1 on 1 mapping for torch/XLA.
|
|
c10::Device backendDeviceToAtenDevice(const BackendDevice& device) {
|
|
return c10::Device(at::kLazy, device.ordinal());
|
|
}
|
|
|
|
c10::optional<BackendDevice> GetBackendDevice(const at::Tensor& tensor) {
|
|
if (auto lt = TryGetLtcTensor(tensor)) {
|
|
return lt->GetDevice();
|
|
}
|
|
return c10::nullopt;
|
|
}
|
|
|
|
c10::optional<BackendDevice> GetBackendDevice() {
|
|
return c10::nullopt;
|
|
}
|
|
|
|
} // namespace lazy
|
|
} // namespace torch
|