wehub-resource-sync
595 lines
26 KiB
C++
595 lines
26 KiB
C++
/* Copyright 2022 The TensorFlow 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.
|
|
==============================================================================*/
|
|
|
|
#include <Python.h>
|
|
|
|
#include <cstdint>
|
|
#include <memory>
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
#include "absl/strings/str_format.h"
|
|
#include "absl/strings/string_view.h"
|
|
#include "pybind11/pybind11.h" // from @pybind11
|
|
#include "pybind11/stl.h" // from @pybind11
|
|
#include "pybind11_abseil/absl_casters.h" // from @pybind11_abseil
|
|
#include "pybind11_protobuf/native_proto_caster.h" // from @pybind11_protobuf
|
|
#include "tensorflow/c/eager/c_api.h"
|
|
#include "tensorflow/core/framework/types.pb.h"
|
|
#include "tensorflow/dtensor/cc/dtensor_device.h"
|
|
#include "tensorflow/dtensor/cc/tensor_layout.h"
|
|
#include "tensorflow/dtensor/proto/layout.pb.h"
|
|
#include "tensorflow/python/eager/pywrap_tensor.h"
|
|
#include "tensorflow/python/eager/pywrap_tfe.h"
|
|
#include "tensorflow/python/lib/core/pybind11_lib.h"
|
|
#include "tensorflow/python/lib/core/pybind11_status.h"
|
|
#include "tensorflow/python/lib/core/safe_pyobject_ptr.h"
|
|
#include "tensorflow/python/util/util.h"
|
|
|
|
namespace py = ::pybind11;
|
|
using tensorflow::dtensor::AddMesh;
|
|
using tensorflow::dtensor::AllocateDTensorDevice;
|
|
using tensorflow::dtensor::ClearTPUCoreIDs;
|
|
using tensorflow::dtensor::ExperimentalClearDefaultLayout;
|
|
using tensorflow::dtensor::ExperimentalClearDefaultMesh;
|
|
using tensorflow::dtensor::ExperimentalSetDefaultLayout;
|
|
using tensorflow::dtensor::ExperimentalSetDefaultMesh;
|
|
using tensorflow::dtensor::FetchLayout;
|
|
using tensorflow::dtensor::GetStats;
|
|
using tensorflow::dtensor::IsDTensor;
|
|
using tensorflow::dtensor::IsSparseDTensor;
|
|
using tensorflow::dtensor::Layout;
|
|
using tensorflow::dtensor::Mesh;
|
|
using tensorflow::dtensor::Pack;
|
|
using tensorflow::dtensor::SetIteratorElementLayouts;
|
|
using tensorflow::dtensor::SetTPUCoreIDs;
|
|
using tensorflow::dtensor::SparsePack;
|
|
using tensorflow::dtensor::TPUCoreIDsToLocations;
|
|
using tensorflow::dtensor::TPUCoreLocationsToIDs;
|
|
using tensorflow::dtensor::Unpack;
|
|
|
|
void PyXDecref(PyObject* obj) { Py_XDECREF(obj); }
|
|
|
|
void CallDelete_Device(PyObject* capsule) {
|
|
delete reinterpret_cast<TFE_CustomDevice*>(
|
|
PyCapsule_GetPointer(capsule, "TFE_CustomDevice"));
|
|
}
|
|
|
|
void CallDelete_DeviceInfo(PyObject* capsule) {
|
|
void (*destructor)(void*) =
|
|
reinterpret_cast<void (*)(void*)>(PyCapsule_GetContext(capsule));
|
|
destructor(PyCapsule_GetPointer(capsule, "TFE_CustomDevice_DeviceInfo"));
|
|
}
|
|
|
|
bool CheckResourceVariable(PyObject* item) {
|
|
if (tensorflow::swig::IsResourceVariable(item)) {
|
|
tensorflow::Safe_PyObjectPtr handle(
|
|
PyObject_GetAttrString(item, "_handle"));
|
|
return EagerTensor_CheckExact(handle.get());
|
|
}
|
|
|
|
return false;
|
|
}
|
|
// Supports 2 cases:
|
|
// i) input is an EagerTensor.
|
|
// ii) input is an arbitrary python list/tuple.
|
|
void ConvertToTensor(TFE_Context* ctx, PyObject* input,
|
|
tensorflow::Safe_PyObjectPtr* output_handle,
|
|
TF_Status* status) {
|
|
if (EagerTensor_CheckExact(input)) {
|
|
// Input is already a EagerTensor so increment the reference, since the
|
|
// caller will use it through output_handle.
|
|
Py_INCREF(input);
|
|
output_handle->reset(input);
|
|
return;
|
|
}
|
|
if (CheckResourceVariable(input)) {
|
|
TF_SetStatus(status, TF_INVALID_ARGUMENT,
|
|
"Variable input is not supported.");
|
|
return;
|
|
}
|
|
TFE_TensorHandle* handle =
|
|
tensorflow::ConvertToEagerTensor(ctx, input, tensorflow::DT_INVALID);
|
|
if (handle == nullptr) {
|
|
TF_SetStatus(status, TF_INTERNAL, "Failure converting to eager tensor.");
|
|
return;
|
|
}
|
|
output_handle->reset(EagerTensorFromHandle(handle));
|
|
}
|
|
|
|
PYBIND11_MODULE(_pywrap_dtensor_device, m) {
|
|
pybind11_protobuf::ImportNativeProtoCasters();
|
|
m.def("Allocate", [](const std::string& name, bool is_async,
|
|
int in_flight_nodes_limit) {
|
|
TFE_CustomDevice* device = new TFE_CustomDevice;
|
|
std::unique_ptr<PyObject, decltype(&PyXDecref)> device_capsule(
|
|
PyCapsule_New(device, "TFE_CustomDevice", &CallDelete_Device),
|
|
PyXDecref);
|
|
void* device_info = nullptr;
|
|
std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
|
|
TF_NewStatus(), TF_DeleteStatus);
|
|
AllocateDTensorDevice(name, device, &device_info, is_async,
|
|
in_flight_nodes_limit, status.get());
|
|
if (TF_GetCode(status.get()) != TF_OK) {
|
|
PyErr_SetString(PyExc_ValueError, TF_Message(status.get()));
|
|
throw py::error_already_set();
|
|
}
|
|
std::unique_ptr<PyObject, decltype(&PyXDecref)> device_info_capsule(
|
|
PyCapsule_New(device_info, "TFE_CustomDevice_DeviceInfo",
|
|
&CallDelete_DeviceInfo),
|
|
PyXDecref);
|
|
// The PyCapsule destructor needs a pointer to the destructor for
|
|
// DeviceInfo.
|
|
PyCapsule_SetContext(device_info_capsule.get(),
|
|
reinterpret_cast<void*>(device->delete_device));
|
|
if (PyErr_Occurred()) throw py::error_already_set();
|
|
return pybind11::reinterpret_steal<pybind11::object>(
|
|
PyTuple_Pack(2, device_capsule.get(), device_info_capsule.get()));
|
|
});
|
|
m.def("AddMesh", [](const py::capsule& device_info,
|
|
const std::string& serialized_mesh, bool is_host_mesh) {
|
|
std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
|
|
TF_NewStatus(), TF_DeleteStatus);
|
|
AddMesh(
|
|
serialized_mesh,
|
|
PyCapsule_GetPointer(device_info.ptr(), "TFE_CustomDevice_DeviceInfo"),
|
|
is_host_mesh, status.get());
|
|
if (TF_GetCode(status.get()) != TF_OK) {
|
|
PyErr_SetString(PyExc_ValueError, TF_Message(status.get()));
|
|
throw py::error_already_set();
|
|
}
|
|
});
|
|
m.def(
|
|
"ExperimentalSetDefaultLayout",
|
|
[](const py::capsule& device_info, const std::string& serialized_layout) {
|
|
std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
|
|
TF_NewStatus(), TF_DeleteStatus);
|
|
ExperimentalSetDefaultLayout(
|
|
serialized_layout,
|
|
PyCapsule_GetPointer(device_info.ptr(),
|
|
"TFE_CustomDevice_DeviceInfo"),
|
|
status.get());
|
|
if (TF_GetCode(status.get()) != TF_OK) {
|
|
PyErr_SetString(PyExc_ValueError, TF_Message(status.get()));
|
|
throw py::error_already_set();
|
|
}
|
|
});
|
|
m.def("ExperimentalClearDefaultLayout", [](const py::capsule& device_info) {
|
|
std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
|
|
TF_NewStatus(), TF_DeleteStatus);
|
|
ExperimentalClearDefaultLayout(
|
|
PyCapsule_GetPointer(device_info.ptr(), "TFE_CustomDevice_DeviceInfo"),
|
|
status.get());
|
|
if (TF_GetCode(status.get()) != TF_OK) {
|
|
PyErr_SetString(PyExc_ValueError, TF_Message(status.get()));
|
|
throw py::error_already_set();
|
|
}
|
|
});
|
|
m.def("ExperimentalSetDefaultMesh", [](const py::capsule& device_info,
|
|
const std::string& serialized_mesh) {
|
|
std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
|
|
TF_NewStatus(), TF_DeleteStatus);
|
|
ExperimentalSetDefaultMesh(
|
|
serialized_mesh,
|
|
PyCapsule_GetPointer(device_info.ptr(), "TFE_CustomDevice_DeviceInfo"),
|
|
status.get());
|
|
if (TF_GetCode(status.get()) != TF_OK) {
|
|
PyErr_SetString(PyExc_ValueError, TF_Message(status.get()));
|
|
throw py::error_already_set();
|
|
}
|
|
});
|
|
m.def("ExperimentalClearDefaultMesh", [](const py::capsule& device_info) {
|
|
std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
|
|
TF_NewStatus(), TF_DeleteStatus);
|
|
ExperimentalClearDefaultMesh(
|
|
PyCapsule_GetPointer(device_info.ptr(), "TFE_CustomDevice_DeviceInfo"),
|
|
status.get());
|
|
if (TF_GetCode(status.get()) != TF_OK) {
|
|
PyErr_SetString(PyExc_ValueError, TF_Message(status.get()));
|
|
throw py::error_already_set();
|
|
}
|
|
});
|
|
m.def("SetTPUCoreIDs", [](const py::capsule& device_info,
|
|
const std::string& mesh_name,
|
|
const std::vector<int>& tpu_core_ids) {
|
|
std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
|
|
TF_NewStatus(), TF_DeleteStatus);
|
|
SetTPUCoreIDs(
|
|
mesh_name, tpu_core_ids,
|
|
PyCapsule_GetPointer(device_info.ptr(), "TFE_CustomDevice_DeviceInfo"),
|
|
status.get());
|
|
if (TF_GetCode(status.get()) != TF_OK) {
|
|
PyErr_SetString(PyExc_ValueError, TF_Message(status.get()));
|
|
throw py::error_already_set();
|
|
}
|
|
});
|
|
m.def("ClearTPUCoreIDs", [](const py::capsule& device_info) {
|
|
ClearTPUCoreIDs(
|
|
PyCapsule_GetPointer(device_info.ptr(), "TFE_CustomDevice_DeviceInfo"));
|
|
});
|
|
m.def("TPUCoreIDsToLocations", [](const py::handle& context,
|
|
const py::capsule& device_info,
|
|
const std::vector<int>& tpu_core_ids) {
|
|
return TPUCoreIDsToLocations(
|
|
static_cast<TFE_Context*>(
|
|
PyCapsule_GetPointer(context.ptr(), "TFE_Context")),
|
|
tpu_core_ids,
|
|
PyCapsule_GetPointer(device_info.ptr(), "TFE_CustomDevice_DeviceInfo"));
|
|
});
|
|
m.def("TPUCoreLocationsToIDs",
|
|
[](const py::handle& context, const py::capsule& device_info,
|
|
const std::vector<std::vector<int>>& tpu_core_locations) {
|
|
return TPUCoreLocationsToIDs(
|
|
static_cast<TFE_Context*>(
|
|
PyCapsule_GetPointer(context.ptr(), "TFE_Context")),
|
|
tpu_core_locations,
|
|
PyCapsule_GetPointer(device_info.ptr(),
|
|
"TFE_CustomDevice_DeviceInfo"));
|
|
});
|
|
m.def("Pack", [](const py::handle& context, const py::handle& input_tensors,
|
|
const std::string& string_layout,
|
|
const py::capsule& device_info, const bool is_sparse) {
|
|
std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
|
|
TF_NewStatus(), TF_DeleteStatus);
|
|
TFE_Context* ctx = static_cast<TFE_Context*>(
|
|
PyCapsule_GetPointer(context.ptr(), "TFE_Context"));
|
|
// Convert each python object to safe py eagertensors.
|
|
std::vector<tensorflow::Safe_PyObjectPtr> py_eager_tensor_handles;
|
|
Py_ssize_t len = PyList_Size(input_tensors.ptr());
|
|
py_eager_tensor_handles.resize(len);
|
|
|
|
for (Py_ssize_t i = 0; i < len; ++i) {
|
|
PyObject* elem = PyList_GetItem(input_tensors.ptr(), i);
|
|
ConvertToTensor(ctx, elem, &py_eager_tensor_handles[i], status.get());
|
|
|
|
if (tensorflow::MaybeRaiseExceptionFromTFStatus(status.get(), nullptr))
|
|
return tensorflow::PyoOrThrow(nullptr);
|
|
}
|
|
std::vector<TFE_TensorHandle*> input_vector;
|
|
input_vector.resize(len);
|
|
for (int i = 0; i < len; ++i)
|
|
input_vector[i] = EagerTensor_Handle(py_eager_tensor_handles[i].get());
|
|
TFE_TensorHandle* packed_tensor;
|
|
if (is_sparse) {
|
|
auto size = input_vector.size() / 3;
|
|
packed_tensor = SparsePack(
|
|
ctx,
|
|
/*num_inputs=*/input_vector.size() / 3,
|
|
/*indices=*/
|
|
std::vector<TFE_TensorHandle*>(input_vector.begin(),
|
|
input_vector.begin() + size)
|
|
.data(),
|
|
/*values=*/
|
|
std::vector<TFE_TensorHandle*>(input_vector.begin() + size,
|
|
input_vector.begin() + 2 * size)
|
|
.data(),
|
|
/*shapes=*/
|
|
std::vector<TFE_TensorHandle*>(input_vector.begin() + 2 * size,
|
|
input_vector.end())
|
|
.data(),
|
|
string_layout, device_info, status.get());
|
|
} else {
|
|
packed_tensor = Pack(ctx, input_vector.size(), input_vector.data(),
|
|
string_layout, device_info, status.get());
|
|
}
|
|
if (tensorflow::MaybeRaiseExceptionFromTFStatus(status.get(), nullptr))
|
|
return tensorflow::PyoOrThrow(nullptr);
|
|
// Convert c++ packed tensor handle into a python eager tensor object.
|
|
tensorflow::Safe_PyObjectPtr flat_result(PyList_New(1));
|
|
PyList_SET_ITEM(flat_result.get(), 0, EagerTensorFromHandle(packed_tensor));
|
|
auto* result = PyList_GET_ITEM(flat_result.get(), 0);
|
|
Py_INCREF(result);
|
|
return tensorflow::PyoOrThrow(result);
|
|
});
|
|
m.def("Unpack", [](const py::handle& context,
|
|
const py::handle& dtensor_handle,
|
|
const py::capsule& device_info) {
|
|
std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
|
|
TF_NewStatus(), TF_DeleteStatus);
|
|
|
|
if (!EagerTensor_CheckExact(dtensor_handle.ptr())) {
|
|
throw py::type_error(absl::StrFormat("Expecting a Tensor, got %s",
|
|
py::str(dtensor_handle.get_type())));
|
|
}
|
|
TFE_TensorHandle* tensor_handle = EagerTensor_Handle(dtensor_handle.ptr());
|
|
std::vector<TFE_TensorHandle*> unpacked_handles =
|
|
Unpack(static_cast<TFE_Context*>(
|
|
PyCapsule_GetPointer(context.ptr(), "TFE_Context")),
|
|
tensor_handle, device_info, status.get());
|
|
|
|
if (tensorflow::MaybeRaiseExceptionFromTFStatus(status.get(), nullptr))
|
|
return tensorflow::PyoOrThrow(nullptr);
|
|
// Convert all TFE_TensorHandles to py EagerTensor and
|
|
// return a python list of them.
|
|
int num_outputs = unpacked_handles.size();
|
|
PyObject* result(PyList_New(num_outputs));
|
|
for (int i = 0; i < num_outputs; ++i) {
|
|
PyList_SET_ITEM(result, i, EagerTensorFromHandle(unpacked_handles[i]));
|
|
}
|
|
return tensorflow::PyoOrThrow(result);
|
|
});
|
|
m.def(
|
|
"FetchLayout",
|
|
[](const py::handle& context, const py::handle& dtensor_handle,
|
|
const py::capsule& device_info) -> py::object {
|
|
std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
|
|
TF_NewStatus(), TF_DeleteStatus);
|
|
if (!EagerTensor_CheckExact(dtensor_handle.ptr())) {
|
|
return py::none();
|
|
}
|
|
TFE_TensorHandle* tensor_handle =
|
|
EagerTensor_Handle(dtensor_handle.ptr());
|
|
std::string layout_string =
|
|
FetchLayout(static_cast<TFE_Context*>(
|
|
PyCapsule_GetPointer(context.ptr(), "TFE_Context")),
|
|
tensor_handle, device_info, status.get());
|
|
if (tensorflow::MaybeRaiseExceptionFromTFStatus(status.get(), nullptr))
|
|
return tensorflow::PyoOrThrow(nullptr);
|
|
return tensorflow::PyoOrThrow(
|
|
PyUnicode_FromString(layout_string.c_str()));
|
|
});
|
|
m.def("IsDTensor", [](const py::handle& context,
|
|
const py::handle& dtensor_handle,
|
|
const py::capsule& device_info) {
|
|
if (!EagerTensor_CheckExact(dtensor_handle.ptr())) {
|
|
return false;
|
|
}
|
|
std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
|
|
TF_NewStatus(), TF_DeleteStatus);
|
|
TFE_TensorHandle* tensor_handle = EagerTensor_Handle(dtensor_handle.ptr());
|
|
bool is_dtensor = IsDTensor(static_cast<TFE_Context*>(PyCapsule_GetPointer(
|
|
context.ptr(), "TFE_Context")),
|
|
tensor_handle, device_info, status.get());
|
|
if (TF_GetCode(status.get()) != TF_OK) {
|
|
PyErr_SetString(PyExc_ValueError, TF_Message(status.get()));
|
|
throw py::error_already_set();
|
|
}
|
|
return is_dtensor;
|
|
});
|
|
m.def("IsSparseDTensor", [](const py::handle& context,
|
|
const py::handle& dtensor_handle,
|
|
const py::capsule& device_info) {
|
|
if (!EagerTensor_CheckExact(dtensor_handle.ptr())) {
|
|
return false;
|
|
}
|
|
std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
|
|
TF_NewStatus(), TF_DeleteStatus);
|
|
|
|
TFE_TensorHandle* tensor_handle = EagerTensor_Handle(dtensor_handle.ptr());
|
|
bool is_sparse =
|
|
IsSparseDTensor(static_cast<TFE_Context*>(
|
|
PyCapsule_GetPointer(context.ptr(), "TFE_Context")),
|
|
tensor_handle, device_info, status.get());
|
|
|
|
if (TF_GetCode(status.get()) != TF_OK) {
|
|
PyErr_SetString(PyExc_ValueError, TF_Message(status.get()));
|
|
throw py::error_already_set();
|
|
}
|
|
return is_sparse;
|
|
});
|
|
m.def("GetStats", [](const py::handle& context,
|
|
const py::capsule& device_info) {
|
|
std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
|
|
TF_NewStatus(), TF_DeleteStatus);
|
|
return GetStats(static_cast<TFE_Context*>(
|
|
PyCapsule_GetPointer(context.ptr(), "TFE_Context")),
|
|
device_info, status.get());
|
|
});
|
|
m.def("SetIteratorElementLayouts",
|
|
[](const py::handle& context, const py::handle& dtensor_handle,
|
|
const std::vector<std::string>& element_layouts,
|
|
const py::capsule& device_info) {
|
|
if (!EagerTensor_CheckExact(dtensor_handle.ptr())) {
|
|
throw py::type_error(
|
|
absl::StrFormat("Expecting a Tensor, got %s",
|
|
py::str(dtensor_handle.get_type())));
|
|
}
|
|
std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
|
|
TF_NewStatus(), TF_DeleteStatus);
|
|
TFE_TensorHandle* tensor_handle =
|
|
EagerTensor_Handle(dtensor_handle.ptr());
|
|
SetIteratorElementLayouts(
|
|
static_cast<TFE_Context*>(
|
|
PyCapsule_GetPointer(context.ptr(), "TFE_Context")),
|
|
tensor_handle, element_layouts, device_info, status.get());
|
|
});
|
|
py::class_<Mesh>(m, "Mesh")
|
|
.def(py::init([](Mesh& mesh) { return mesh; }), py::arg("mesh"),
|
|
"Create a copy of a mesh.")
|
|
.def(py::init(&Mesh::CreateMesh))
|
|
.def(py::init([](absl::string_view single_device) {
|
|
auto mesh = Mesh::GetSingleDeviceMesh(single_device);
|
|
if (!mesh.ok()) {
|
|
throw py::value_error(std::string(mesh.status().message()));
|
|
}
|
|
return *mesh;
|
|
}),
|
|
py::arg("single_device"), "Creates a single device mesh.")
|
|
.def(py::init([](const tensorflow::dtensor::MeshProto& proto) {
|
|
auto mesh = Mesh::ParseFromProto(proto);
|
|
if (!mesh.ok()) {
|
|
throw py::value_error(std::string(mesh.status().message()));
|
|
}
|
|
return *mesh;
|
|
}),
|
|
py::arg("mesh_proto"), "Returns a Mesh from a MeshProto.")
|
|
.def(py::init([](absl::string_view mesh_str) {
|
|
auto mesh = Mesh::FromString(mesh_str);
|
|
if (!mesh.ok()) {
|
|
throw py::value_error(std::string(mesh.status().message()));
|
|
}
|
|
return *mesh;
|
|
}),
|
|
py::arg("mesh_str"), "Returns a Mesh from a string.")
|
|
.def_property_readonly("name", &Mesh::name)
|
|
.def_property_readonly("dim_names", &Mesh::MeshDimNames)
|
|
.def_property_readonly("size", &Mesh::num_devices)
|
|
.def_property_readonly("single_device", &Mesh::single_device)
|
|
.def("__contains__", &Mesh::IsMeshDim, py::arg("dim_name"))
|
|
.def("__eq__", &Mesh::operator==)
|
|
.def("to_string", &Mesh::ToString,
|
|
"Returns string representation of Mesh.")
|
|
.def("is_single_device", &Mesh::IsSingleDevice,
|
|
"Returns True if the mesh represents a non-distributed device.")
|
|
.def("contains_dim", &Mesh::IsMeshDim, py::arg("dim_name"),
|
|
"Returns True if a Mesh contains the given dimension name.")
|
|
.def(
|
|
"dim_size",
|
|
[](const Mesh& mesh, absl::string_view name) {
|
|
auto dim_size = mesh.dim_size(name);
|
|
if (!dim_size.ok()) {
|
|
throw py::value_error(std::string(dim_size.status().message()));
|
|
}
|
|
return *dim_size;
|
|
},
|
|
py::arg("dim_name"), "Returns the size of mesh dimension.")
|
|
.def("device_type", &Mesh::device_type,
|
|
"Returns the device_type of a Mesh.")
|
|
.def("host_mesh", &Mesh::host_mesh,
|
|
"Returns a host mesh corresponding to this mesh.")
|
|
.def("num_local_devices", &Mesh::num_local_devices,
|
|
"Returns the number of local devices.")
|
|
.def("min_global_device_id", &Mesh::min_global_device_id,
|
|
"Returns the minimum global device ID.")
|
|
.def("is_remote", &Mesh::is_remote,
|
|
"Returns True if a Mesh contains only remote devices.")
|
|
.def("local_device_ids", &Mesh::local_device_ids,
|
|
"Returns a list of local device IDs.")
|
|
.def("local_devices", &Mesh::local_devices,
|
|
"Returns a list of local device specs "
|
|
"represented as strings.")
|
|
.def("global_device_ids", &Mesh::global_device_ids,
|
|
"Returns a list of global device IDs.")
|
|
.def("global_devices", &Mesh::global_devices,
|
|
"Returns a list of global device specs "
|
|
"represented as strings.")
|
|
.def("shape", &Mesh::dim_sizes, "Returns the shape of the mesh.")
|
|
.def("use_xla_spmd", &Mesh::use_xla_spmd,
|
|
"Returns True if Mesh will use XLA for SPMD "
|
|
"instead of DTensor SPMD.")
|
|
.def(
|
|
"as_proto",
|
|
[](const Mesh& mesh) {
|
|
auto mesh_proto = mesh.ToProto();
|
|
if (!mesh_proto.ok()) {
|
|
throw py::value_error(std::string(mesh_proto.status().message()));
|
|
}
|
|
return *mesh_proto;
|
|
},
|
|
"Returns the MeshProto protobuf message.")
|
|
.def("device_location", [](const Mesh& mesh, int device_id) {
|
|
auto location = mesh.device_location(device_id);
|
|
if (!location.ok()) {
|
|
throw py::value_error(std::string(location.status().message()));
|
|
}
|
|
return std::vector<int64_t>(location->begin(), location->end());
|
|
});
|
|
|
|
py::enum_<Layout::LayoutType>(m, "LayoutType")
|
|
.value("STATIC", Layout::LayoutType::kStatic)
|
|
.value("PARTED", Layout::LayoutType::kParted)
|
|
.value("SINGLE_DEVICE", Layout::LayoutType::kSingleDevice);
|
|
|
|
py::class_<Layout>(m, "Layout")
|
|
.def(py::init([](Layout& layout) { return layout; }), py::arg("layout"),
|
|
"Create a copy of a layout.")
|
|
.def(py::init([](Layout::LayoutType type,
|
|
const std::vector<std::string>& sharding_specs,
|
|
const Mesh& mesh) {
|
|
auto layout = Layout::GetLayout(type, sharding_specs, mesh);
|
|
if (!layout.ok()) {
|
|
throw py::value_error(std::string(layout.status().message()));
|
|
}
|
|
return *layout;
|
|
}),
|
|
py::arg("type"), py::arg("sharding_specs"), py::arg("mesh"))
|
|
.def(py::init([](const tensorflow::dtensor::LayoutProto& proto) {
|
|
auto layout = Layout::FromProto(proto);
|
|
if (!layout.ok()) {
|
|
throw py::value_error(std::string(layout.status().message()));
|
|
}
|
|
return *layout;
|
|
}),
|
|
py::arg("layout_proto"), "Returns a Layout from a LayoutProto.")
|
|
.def(py::init([](absl::string_view layout_str) {
|
|
auto layout = Layout::FromString(layout_str);
|
|
if (!layout.ok()) {
|
|
throw py::value_error(std::string(layout.status().message()));
|
|
}
|
|
return *layout;
|
|
}),
|
|
py::arg("layout_str"), "Returns a Layout from a string.")
|
|
.def(py::init(&Layout::ReplicatedOnMesh), py::arg("mesh"),
|
|
py::arg("rank"), "Returns a replicated layout.")
|
|
.def(py::init(&Layout::BatchShardedOnMesh), py::arg("mesh"),
|
|
py::arg("rank"), py::arg("batch_dim"), py::arg("axis"),
|
|
"Returns a batch sharded layout.")
|
|
.def(py::init([](const Mesh& mesh) {
|
|
auto layout = Layout::GetSingleDeviceLayout(mesh);
|
|
if (!layout.ok()) {
|
|
throw py::value_error(std::string(layout.status().message()));
|
|
}
|
|
return *layout;
|
|
}),
|
|
py::arg("mesh"), "Returns a single device layout.")
|
|
.def("__eq__", &Layout::operator==)
|
|
.def(
|
|
"as_proto",
|
|
[](const Layout& layout) {
|
|
auto layout_proto = layout.ToProto();
|
|
if (!layout_proto.ok()) {
|
|
throw py::value_error(
|
|
std::string(layout_proto.status().message()));
|
|
}
|
|
return *layout_proto;
|
|
},
|
|
"Returns the LayoutProto protobuf message.")
|
|
.def("to_string", &Layout::ToString)
|
|
.def("to_parted", &Layout::ToParted)
|
|
.def_property_readonly("type", &Layout::type)
|
|
.def_property_readonly("sharding_specs", &Layout::sharding_spec_strs)
|
|
.def_property_readonly("rank", &Layout::rank)
|
|
.def_property_readonly("mesh", &Layout::mesh)
|
|
.def("is_fully_replicated", &Layout::IsFullyReplicated,
|
|
"Returns True if all tensor axes are replicated.")
|
|
.def("is_batch_parallel",
|
|
[](const Layout& layout) { return layout.IsBatchParallel(); })
|
|
.def("is_single_device", &Layout::IsSingleDevice,
|
|
"Returns True if the Layout represents a non-distributed device.")
|
|
.def(
|
|
"num_shards",
|
|
[](const Layout& layout, int dim) {
|
|
return layout.num_shards_for_dim(dim);
|
|
},
|
|
py::arg("idx"),
|
|
"Returns the number of shards for tensor dimension `idx`.")
|
|
.def(
|
|
"global_shape_from_local_shape",
|
|
[](const Layout& layout, std::vector<int64_t> local_shape) {
|
|
return py::tuple(
|
|
py::cast(layout.GlobalShapeFromLocalShape(local_shape)));
|
|
},
|
|
py::arg("local_shape"),
|
|
"Returns the global shape computed from this local shape.")
|
|
.def(
|
|
"local_shape_from_global_shape",
|
|
[](const Layout& layout, std::vector<int64_t> global_shape) {
|
|
return py::tuple(
|
|
py::cast(layout.LocalShapeFromGlobalShape(global_shape)));
|
|
},
|
|
py::arg("global_shape"),
|
|
"Returns the local shape computed from this global shape.");
|
|
}
|