wehub-resource-sync
113 lines
4.0 KiB
C++
113 lines
4.0 KiB
C++
/* Copyright 2023 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.
|
|
==============================================================================*/
|
|
|
|
#ifndef TENSORFLOW_COMPILER_JIT_PJRT_BASE_DEVICE_H_
|
|
#define TENSORFLOW_COMPILER_JIT_PJRT_BASE_DEVICE_H_
|
|
|
|
#include <string>
|
|
#include <utility>
|
|
#include <vector>
|
|
|
|
#include "tensorflow/compiler/tf2xla/layout_util.h"
|
|
#include "tensorflow/core/common_runtime/local_device.h"
|
|
#include "tensorflow/core/framework/device_base.h"
|
|
|
|
namespace tensorflow {
|
|
|
|
// tensorflow::PjRtBaseDevice replaces the deprecated tensorflow::XlaDevice.
|
|
// This accelerator agnostic device is mainly used to store metadata.
|
|
class PjRtBaseDevice : public LocalDevice {
|
|
public:
|
|
// Stores metadata about the PjRtBaseDevice.
|
|
class Metadata {
|
|
public:
|
|
Metadata(const DeviceType& jit_device_type,
|
|
std::vector<XlaShapeLayoutHelpers::ShapeDeterminationFns>
|
|
shape_determination_fns)
|
|
: jit_device_type_(jit_device_type),
|
|
shape_determination_fns_(std::move(shape_determination_fns)) {}
|
|
|
|
// The index of the device on this host.
|
|
int device_ordinal() const;
|
|
|
|
const DeviceType& jit_device_type() const { return jit_device_type_; }
|
|
const XlaShapeLayoutHelpers::ShapeDeterminationFns&
|
|
default_shape_determination_fns() const {
|
|
return shape_determination_fns_.at(0);
|
|
}
|
|
|
|
const XlaShapeLayoutHelpers::ShapeDeterminationFns&
|
|
shape_determination_fns_at(int i) const {
|
|
return shape_determination_fns_[i];
|
|
}
|
|
|
|
private:
|
|
const DeviceType jit_device_type_;
|
|
std::vector<XlaShapeLayoutHelpers::ShapeDeterminationFns>
|
|
shape_determination_fns_;
|
|
|
|
Metadata(const Metadata&) = delete;
|
|
void operator=(const Metadata&) = delete;
|
|
};
|
|
|
|
struct Options {
|
|
// The device name's prefix (e.g., "/task:7")
|
|
std::string device_name_prefix;
|
|
|
|
// The name of the device (e.g., "TPU")
|
|
std::string device_name;
|
|
|
|
// The index of the device.
|
|
int device_ordinal = -1;
|
|
|
|
// The name of the compilation device, also referred to as jit_device_type.
|
|
// (e.g., "XLA_CPU_JIT");
|
|
std::string compilation_device_name;
|
|
|
|
// A vector of ShapeDeterminationFn (i.e., a bundle of LayoutSelectionFn,
|
|
// ShapeRepresentationFn). Each bundle describes how the on-host shapes of
|
|
// a) argument and return value, for entry computations b) variables, for
|
|
// all computations, should be represented in XLA. Parameters/return values
|
|
// will be shaped according to the function pair, and reshaped back to/from
|
|
// their declared shapes for computations. Must be non-empty.
|
|
std::vector<XlaShapeLayoutHelpers::ShapeDeterminationFns>
|
|
shape_determination_fns;
|
|
|
|
Options(std::string device_name_prefix, std::string device_name,
|
|
int device_ordinal, std::string compilation_device_name,
|
|
std::vector<XlaShapeLayoutHelpers::ShapeDeterminationFns>
|
|
shape_determination_fns)
|
|
: device_name_prefix(device_name_prefix),
|
|
device_name(device_name),
|
|
device_ordinal(device_ordinal),
|
|
compilation_device_name(compilation_device_name),
|
|
shape_determination_fns(shape_determination_fns) {}
|
|
};
|
|
|
|
// Creates a new PJRT base device.
|
|
PjRtBaseDevice(const SessionOptions& session_options, const Options& options);
|
|
|
|
static absl::StatusOr<const PjRtBaseDevice::Metadata*> GetMetadataFromDevice(
|
|
DeviceBase* device);
|
|
|
|
private:
|
|
// The metadata of this PjRtBaseDevice.
|
|
const Metadata metadata_;
|
|
};
|
|
|
|
} // namespace tensorflow
|
|
|
|
#endif // TENSORFLOW_COMPILER_JIT_PJRT_BASE_DEVICE_H_
|