Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Default Device Ordinal API #1079

Merged
merged 2 commits into from
Jul 19, 2022
Merged

Default Device Ordinal API #1079

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
a9443ea
Add default device ordinal API
antoniojkim Jul 15, 2022
45ec643
Fix reference backend
antoniojkim Jul 16, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions python/torch_mlir/csrc/base_lazy_backend/backend_impl.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -192,5 +192,13 @@ BackendDevice TorchMlirBackendImpl::GetBackendDevice(c10::Device device) const {
return BackendDevice(GetDefaultDeviceType(), device.index());
}

int64_t TorchMlirBackendImpl::GetDefaultDeviceOrdinal() const {
return default_device_ordinal;
}

void TorchMlirBackendImpl::SetDefaultDeviceOrdinal(int64_t ordinal) {
default_device_ordinal = ordinal;
}

} // namespace lazy
} // namespace torch
7 changes: 7 additions & 0 deletions python/torch_mlir/csrc/base_lazy_backend/backend_impl.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -153,6 +153,10 @@ class TORCH_API TorchMlirBackendImpl : public BackendImplInterface {
// identity mappings.
virtual BackendDevice GetBackendDevice(c10::Device device) const override;

virtual int64_t GetDefaultDeviceOrdinal() const override;

virtual void SetDefaultDeviceOrdinal(int64_t ordinal) override;

/**
* Debug/Metrics
* */
Expand All @@ -164,6 +168,9 @@ class TORCH_API TorchMlirBackendImpl : public BackendImplInterface {
// virtual std::string GetComputationBackendText(
// const ComputationPtr computation
// ) const = 0;

protected:
int64_t default_device_ordinal = 0;
};

} // namespace lazy
Expand Down
14 changes: 9 additions & 5 deletions python/torch_mlir/csrc/reference_lazy_backend/backend_impl.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
//
//===----------------------------------------------------------------------===//

#include <c10/core/DeviceType.h>
#include <torch/csrc/lazy/backend/backend_data.h>
#include <torch/csrc/lazy/backend/backend_device.h>
#include <torch/csrc/lazy/backend/lowering_context.h>
Expand All @@ -26,17 +27,19 @@ namespace torch {
namespace lazy {

struct ReferenceLazyBackendDeviceType : public BackendDeviceType {
ReferenceLazyBackendDeviceType(std::string device_type)
ReferenceLazyBackendDeviceType(c10::DeviceType device_type)
: device_type_(device_type) {}
ReferenceLazyBackendDeviceType(int8_t device_type)
: device_type_(static_cast<c10::DeviceType>(device_type)) {}

std::string toString() const override { return device_type_; }
std::string toString() const override { return c10::DeviceTypeName(device_type_); }

std::string device_type_;
c10::DeviceType device_type_;
};

class ReferenceLazyBackendImpl : public torch::lazy::TorchMlirBackendImpl {
public:
ReferenceLazyBackendImpl() : default_device_type_("Magic") {}
ReferenceLazyBackendImpl() : default_device_type_(c10::DeviceType::Lazy) {}

/**
* Configuration
Expand Down Expand Up @@ -128,7 +131,8 @@ class ReferenceLazyBackendImpl : public torch::lazy::TorchMlirBackendImpl {
return std::make_shared<BackendDeviceType>(default_device_type_);
}

void SetDefaultDeviceType(std::string device_type) {

void SetDefaultDeviceType(int8_t device_type) override {
default_device_type_ = ReferenceLazyBackendDeviceType(device_type);
}

Expand Down