wehub-resource-sync
326 lines
12 KiB
C++
326 lines
12 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 "tensorflow/dtensor/mlir/layout_parsing.h"
|
|
|
|
#include <optional>
|
|
#include <string>
|
|
#include <utility>
|
|
#include <vector>
|
|
|
|
#include "absl/types/optional.h"
|
|
#include "absl/types/span.h"
|
|
#include "llvm/ADT/STLExtras.h"
|
|
#include "llvm/Support/Casting.h"
|
|
#include "llvm/Support/FormatVariadic.h"
|
|
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
|
|
#include "mlir/IR/Attributes.h" // from @llvm-project
|
|
#include "mlir/IR/Builders.h" // from @llvm-project
|
|
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
|
|
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
|
|
#include "mlir/IR/Operation.h" // from @llvm-project
|
|
#include "mlir/IR/OperationSupport.h" // from @llvm-project
|
|
#include "mlir/IR/Value.h" // from @llvm-project
|
|
#include "mlir/Support/LLVM.h" // from @llvm-project
|
|
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"
|
|
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
|
|
#include "tensorflow/core/platform/errors.h"
|
|
#include "tensorflow/dtensor/cc/constants.h"
|
|
#include "tensorflow/dtensor/cc/dstatus.h"
|
|
#include "tensorflow/dtensor/cc/tensor_layout.h"
|
|
#include "tensorflow/dtensor/mlir/ir/tf_dtensor.h"
|
|
|
|
namespace tensorflow {
|
|
namespace dtensor {
|
|
namespace {
|
|
|
|
bool OpUsesV2LayoutAnnotation(mlir::Operation* op) {
|
|
return !op->getUsers().empty() &&
|
|
llvm::all_of(op->getUsers(), [](mlir::Operation* user_op) {
|
|
return llvm::isa<mlir::TF::DTensorLayout>(user_op);
|
|
});
|
|
}
|
|
|
|
} // namespace
|
|
|
|
StatusOr<std::optional<Layout>> ExtractSingleLayoutFromOp(
|
|
mlir::Operation* op, std::string attr_name) {
|
|
std::optional<Layout> out;
|
|
|
|
// If v2 layout propagation algorithm is used, parse layout from DTensorLayout
|
|
// op.
|
|
if (OpUsesV2LayoutAnnotation(op)) {
|
|
// If DTensorLayout is used, then DTensorLayout op is the only consumer for
|
|
// the operation output value.
|
|
auto users = op->getUsers();
|
|
out.emplace(
|
|
llvm::cast<mlir::TF::DTensorLayout>(*users.begin()).getLayout());
|
|
} else {
|
|
TF_ASSIGN_OR_RETURN(auto layouts, ExtractLayoutFromOp(op, attr_name));
|
|
if (layouts.empty()) return out;
|
|
if (layouts.size() != 1) {
|
|
return absl::InternalError(absl::StrCat(
|
|
"Extracting single layout on Op that has multiple layout attached is "
|
|
"ambiguous. op : ",
|
|
op->getName().getStringRef().str()));
|
|
}
|
|
out.swap(layouts[0]);
|
|
}
|
|
return out;
|
|
}
|
|
|
|
StatusOr<std::optional<Layout>> ExtractSingleLayoutFromOp(mlir::Operation* op) {
|
|
return ExtractSingleLayoutFromOp(op, kLayoutAttr);
|
|
}
|
|
|
|
StatusOr<Layout> ExtractRequiredSingleLayoutFromOp(mlir::Operation* op) {
|
|
TF_ASSIGN_OR_RETURN(std::optional<Layout> layout,
|
|
ExtractSingleLayoutFromOp(op));
|
|
if (!layout) return absl::InternalError("expected layout missing");
|
|
|
|
return *layout;
|
|
}
|
|
|
|
StatusOr<std::vector<std::optional<Layout>>> ExtractLayoutFromOp(
|
|
mlir::Operation* op, std::string attr_name) {
|
|
std::vector<std::optional<Layout>> outs;
|
|
outs.reserve(op->getNumResults());
|
|
|
|
// If v2 layout propagation algorithm is used, parse layout from DTensorLayout
|
|
// op.
|
|
if (OpUsesV2LayoutAnnotation(op)) {
|
|
for (auto op_result : op->getOpResults()) {
|
|
outs.emplace_back(
|
|
llvm::cast<mlir::TF::DTensorLayout>(*op_result.getUsers().begin())
|
|
.getLayout());
|
|
}
|
|
} else {
|
|
auto serialized_layouts = op->getAttrOfType<mlir::ArrayAttr>(attr_name);
|
|
if (!serialized_layouts) return outs;
|
|
|
|
for (auto const& attr : serialized_layouts) {
|
|
auto attr_str = mlir::cast<mlir::StringAttr>(attr).getValue().str();
|
|
if (!attr_str.empty()) {
|
|
TF_ASSIGN_OR_RETURN(auto layout, Layout::FromString(attr_str));
|
|
outs.emplace_back(std::move(layout));
|
|
} else {
|
|
outs.emplace_back(std::nullopt);
|
|
}
|
|
}
|
|
}
|
|
return outs;
|
|
}
|
|
|
|
StatusOr<std::vector<std::optional<Layout>>> ExtractLayoutFromOp(
|
|
mlir::Operation* op) {
|
|
return ExtractLayoutFromOp(op, kLayoutAttr);
|
|
}
|
|
|
|
StatusOr<std::vector<Layout>> ExtractRequiredLayoutFromOp(mlir::Operation* op) {
|
|
TF_ASSIGN_OR_RETURN(std::vector<std::optional<Layout>> optional_layouts,
|
|
ExtractLayoutFromOp(op));
|
|
std::vector<Layout> layouts;
|
|
for (const absl::optional<Layout>& layout : optional_layouts) {
|
|
if (!layout) return absl::InternalError("expected layout missing");
|
|
layouts.emplace_back(*layout);
|
|
}
|
|
|
|
return layouts;
|
|
}
|
|
|
|
StatusOr<Mesh> ExtractDeviceMeshEnclosingCluster(mlir::Operation* op) {
|
|
auto enclosing_cluster = op->getParentOfType<mlir::tf_device::ClusterOp>();
|
|
if (!enclosing_cluster)
|
|
return absl::InvalidArgumentError(
|
|
"op is not inside a device mesh cluster.");
|
|
|
|
TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshFromOp(enclosing_cluster));
|
|
if (!mesh)
|
|
return absl::InvalidArgumentError(
|
|
"op's enclosing device cluster does not have mesh defined.");
|
|
|
|
return *mesh;
|
|
}
|
|
|
|
StatusOr<std::optional<Mesh>> ExtractDeviceMeshFromOp(mlir::Operation* op) {
|
|
std::optional<Mesh> extracted_mesh;
|
|
if (op == nullptr) return extracted_mesh;
|
|
|
|
auto mesh_str_attr = op->getAttrOfType<mlir::StringAttr>(kMeshAttr);
|
|
if (!mesh_str_attr) return extracted_mesh;
|
|
|
|
TF_ASSIGN_OR_RETURN(Mesh mesh,
|
|
Mesh::FromString(mesh_str_attr.getValue().str()));
|
|
|
|
extracted_mesh.emplace(std::move(mesh));
|
|
return extracted_mesh;
|
|
}
|
|
|
|
StatusOr<std::optional<Layout>> ExtractLayoutFromOperand(mlir::Value operand) {
|
|
if (auto op_result = mlir::dyn_cast<mlir::OpResult>(operand)) {
|
|
mlir::Operation* op = op_result.getDefiningOp();
|
|
std::optional<Layout> out;
|
|
if (auto layout_op = llvm::dyn_cast<mlir::TF::DTensorLayout>(op)) {
|
|
out.emplace(layout_op.getLayout());
|
|
} else {
|
|
const int result_number = op_result.getResultNumber();
|
|
TF_ASSIGN_OR_RETURN(auto layouts, ExtractLayoutFromOp(op, kLayoutAttr));
|
|
|
|
if (layouts.empty()) return out;
|
|
|
|
if (result_number >= layouts.size()) {
|
|
return absl::InternalError(absl::StrCat(
|
|
"Expect to extract the ", result_number,
|
|
"-th output's layout, but "
|
|
"only see ",
|
|
layouts.size(), " outputs: ", op->getName().getStringRef().str()));
|
|
}
|
|
out.swap(layouts[result_number]);
|
|
}
|
|
return out;
|
|
}
|
|
|
|
auto block_arg = mlir::dyn_cast<mlir::BlockArgument>(operand);
|
|
if (!block_arg)
|
|
return absl::InternalError(
|
|
"Operand is not either a OpResult or a BlockArgument. This should not "
|
|
"happen.");
|
|
auto func_op = mlir::dyn_cast_or_null<mlir::func::FuncOp>(
|
|
block_arg.getOwner()->getParentOp());
|
|
if (!func_op) {
|
|
return absl::InvalidArgumentError("op must be enclosed by a function");
|
|
}
|
|
|
|
std::optional<Layout> extracted_layout;
|
|
auto layout_attr = func_op.getArgAttrOfType<mlir::StringAttr>(
|
|
block_arg.getArgNumber(), kCustomDeviceAttr);
|
|
if (!layout_attr) return extracted_layout;
|
|
|
|
TF_ASSIGN_OR_RETURN(auto layout,
|
|
Layout::FromString(layout_attr.getValue().str()));
|
|
extracted_layout.emplace(std::move(layout));
|
|
return extracted_layout;
|
|
}
|
|
|
|
StatusOr<Layout> ExtractRequiredLayoutFromOperand(mlir::Value operand) {
|
|
TF_ASSIGN_OR_RETURN(std::optional<Layout> layout,
|
|
ExtractLayoutFromOperand(operand));
|
|
if (!layout) return absl::InternalError("expected layout missing");
|
|
|
|
return *layout;
|
|
}
|
|
|
|
StatusOr<std::vector<Layout>> ExtractRequiredLayoutFromOperands(
|
|
mlir::Operation* op) {
|
|
std::vector<Layout> layouts;
|
|
for (const auto& operand : op->getOpOperands()) {
|
|
TF_ASSIGN_OR_RETURN(auto operand_layout,
|
|
ExtractRequiredLayoutFromOperand(operand.get()));
|
|
layouts.emplace_back(operand_layout);
|
|
}
|
|
return layouts;
|
|
}
|
|
|
|
void SetLayoutOnOp(mlir::Operation* op, mlir::OpBuilder builder,
|
|
absl::Span<const absl::optional<Layout>> layouts) {
|
|
llvm::SmallVector<std::string, 8> serialized_layouts;
|
|
for (auto const& layout : layouts) {
|
|
serialized_layouts.emplace_back(layout.has_value() ? layout->ToString()
|
|
: "");
|
|
}
|
|
op->setAttr(kLayoutAttr,
|
|
builder.getStrArrayAttr(llvm::SmallVector<llvm::StringRef, 8>(
|
|
serialized_layouts.begin(), serialized_layouts.end())));
|
|
}
|
|
|
|
void SetLayoutOnOp(mlir::Operation* op,
|
|
absl::Span<const absl::optional<Layout>> layouts) {
|
|
SetLayoutOnOp(op, mlir::OpBuilder(op), layouts);
|
|
}
|
|
|
|
void SetSingleLayoutOnOp(mlir::Operation* op, const Layout& layout) {
|
|
SetLayoutOnOp(op, mlir::OpBuilder(op), {std::optional<Layout>(layout)});
|
|
}
|
|
|
|
StatusOr<std::optional<Layout>> ExtractLayoutFromFunctionReturnAttr(
|
|
mlir::func::ReturnOp return_op, const int return_index) {
|
|
std::optional<Layout> layout;
|
|
// If value feeds into func op return op, then check to see if layout
|
|
// attribute is set for the return value.
|
|
auto function = return_op->getParentOfType<mlir::func::FuncOp>();
|
|
auto layout_attr_from_func_result =
|
|
function.getResultAttrOfType<mlir::StringAttr>(return_index,
|
|
kCustomDefaultLayoutAttr);
|
|
if (!layout_attr_from_func_result) return layout;
|
|
|
|
const std::string layout_string =
|
|
layout_attr_from_func_result.getValue().str();
|
|
auto result_layout_or_status = Layout::FromString(layout_string);
|
|
if (!result_layout_or_status.ok())
|
|
return absl::InvalidArgumentError(
|
|
llvm::formatv("Malformed default return layout received. {0} Received "
|
|
"layout : {1}",
|
|
result_layout_or_status.status().message(), layout_string)
|
|
.str());
|
|
|
|
layout.emplace(result_layout_or_status.value());
|
|
return layout;
|
|
}
|
|
|
|
StatusOr<llvm::SmallVector<Layout, 4>> ExtractElementLayoutsFromOperand(
|
|
mlir::OpOperand& input_value) {
|
|
const int operand_index = input_value.getOperandNumber();
|
|
auto defining_op = input_value.get().getDefiningOp();
|
|
|
|
if (defining_op) {
|
|
if (mlir::isa<mlir::TF::DTensorLayout,
|
|
mlir::TF::IteratorGetNextAsOptionalOp>(defining_op)) {
|
|
return ExtractElementLayoutsFromOperand(defining_op->getOpOperand(0));
|
|
}
|
|
}
|
|
|
|
// If we reach this point, we're working with a function argument.
|
|
mlir::Operation* op = input_value.getOwner();
|
|
auto enclosing_function = op->getParentOfType<mlir::func::FuncOp>();
|
|
if (!enclosing_function)
|
|
return absl::InvalidArgumentError(
|
|
llvm::formatv("Could not find iterator at {0}-th input to op: {1}",
|
|
operand_index, op->getName())
|
|
.str());
|
|
|
|
auto block_arg = mlir::dyn_cast<mlir::BlockArgument>(input_value.get());
|
|
auto array_attr = enclosing_function.getArgAttrOfType<mlir::ArrayAttr>(
|
|
block_arg.getArgNumber(), kIteratorElementLayouts);
|
|
if (!array_attr)
|
|
return absl::InvalidArgumentError(
|
|
llvm::formatv(
|
|
"Could not find `{0}` attribute of {1}-th input to op: {2}",
|
|
kIteratorElementLayouts, operand_index, op->getName())
|
|
.str());
|
|
|
|
llvm::SmallVector<Layout, 4> layouts(array_attr.size());
|
|
for (int i = 0; i < array_attr.size(); ++i) {
|
|
layouts[i] =
|
|
Layout::FromString(
|
|
mlir::cast<mlir::StringAttr>(array_attr[i]).getValue().str())
|
|
.value();
|
|
}
|
|
|
|
return layouts;
|
|
}
|
|
|
|
} // namespace dtensor
|
|
} // namespace tensorflow
|