[MLIR:TF/XLA] Side effect analysis.
An analysis that infers control dependencies based on side-effects on known and unknown resources. Side-effecting ops on uknown resources are conservatively treated as interferencing all known resource op accesses. It distinguishes accesses based on whether they are read-only, and read-only ops do not interfer with each other. PiperOrigin-RevId: 282815679 Change-Id: I4f430c5b6cbfb02284c150c85b41f0b81458b6e9
This commit is contained in:
Yuanzhong Xu
TensorFlower Gardener
parent
bbd8a3e0e4
commit
fe77f79256
@ -217,6 +217,7 @@ cc_library(
|
||||
"transforms/shape_inference.cc",
|
||||
"transforms/shape_inference_pass.cc",
|
||||
"transforms/sink_constant.cc",
|
||||
"transforms/test_side_effect_analysis.cc",
|
||||
"transforms/tpu_cluster_formation.cc",
|
||||
"transforms/tpu_merge_variables_with_execute.cc",
|
||||
"transforms/tpu_rewrite_pass.cc",
|
||||
@ -239,6 +240,7 @@ cc_library(
|
||||
":error_util",
|
||||
":export_tf_dialect_op",
|
||||
":mangling_util",
|
||||
":side_effect_analysis",
|
||||
":tensorflow",
|
||||
":tensorflow_optimize_inc_gen",
|
||||
":tpu_rewrite_device_util",
|
||||
@ -981,3 +983,19 @@ cc_library(
|
||||
"@local_config_mlir//:Pass",
|
||||
],
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "side_effect_analysis",
|
||||
srcs = ["analysis/side_effect_analysis.cc"],
|
||||
hdrs = ["analysis/side_effect_analysis.h"],
|
||||
deps = [
|
||||
":tensorflow",
|
||||
"//tensorflow/compiler/tf2xla:resource_operation_table",
|
||||
"//tensorflow/core:framework",
|
||||
"@com_google_absl//absl/strings",
|
||||
"@llvm//:support",
|
||||
"@local_config_mlir//:IR",
|
||||
"@local_config_mlir//:Pass",
|
||||
"@local_config_mlir//:Support",
|
||||
],
|
||||
)
|
||||
|
||||
@ -0,0 +1,374 @@
|
||||
/* Copyright 2019 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/compiler/mlir/tensorflow/analysis/side_effect_analysis.h"
|
||||
|
||||
#include <cstdint>
|
||||
#include <initializer_list>
|
||||
|
||||
#include "absl/strings/str_cat.h"
|
||||
#include "llvm/ADT/DenseMap.h"
|
||||
#include "llvm/ADT/DenseSet.h"
|
||||
#include "llvm/ADT/STLExtras.h"
|
||||
#include "llvm/ADT/SmallVector.h"
|
||||
#include "llvm/ADT/iterator_range.h"
|
||||
#include "llvm/Support/Casting.h"
|
||||
#include "llvm/Support/Debug.h"
|
||||
#include "mlir/IR/Attributes.h" // TF:local_config_mlir
|
||||
#include "mlir/IR/Block.h" // TF:local_config_mlir
|
||||
#include "mlir/IR/Builders.h" // TF:local_config_mlir
|
||||
#include "mlir/IR/Location.h" // TF:local_config_mlir
|
||||
#include "mlir/IR/Operation.h" // TF:local_config_mlir
|
||||
#include "mlir/IR/StandardTypes.h" // TF:local_config_mlir
|
||||
#include "mlir/Support/LLVM.h" // TF:local_config_mlir
|
||||
#include "mlir/Support/LogicalResult.h" // TF:local_config_mlir
|
||||
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
|
||||
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"
|
||||
#include "tensorflow/compiler/tf2xla/resource_operation_table.h"
|
||||
#include "tensorflow/core/framework/resource_mgr.h"
|
||||
|
||||
namespace mlir {
|
||||
namespace TF {
|
||||
|
||||
namespace {
|
||||
|
||||
constexpr int64_t kUnknownResourceId = -1;
|
||||
|
||||
// Returns if a VarHandleOp is anonymous, which means it always creates a new
|
||||
// variable.
|
||||
bool IsResourceHandleAnonymous(TF::VarHandleOp handle) {
|
||||
return handle.shared_name() == tensorflow::ResourceHandle::ANONYMOUS_NAME;
|
||||
}
|
||||
|
||||
// Returns a string unique identifier for a non-anonymous VarHandleOp.
|
||||
std::string GetVarHandleStringId(TF::VarHandleOp handle) {
|
||||
auto device = handle.getAttrOfType<StringAttr>("device");
|
||||
return absl::StrCat(handle.container().str(), "/", handle.shared_name().str(),
|
||||
"/", device ? device.getValue().str() : std::string(""));
|
||||
}
|
||||
|
||||
// Finds a unique ID for a VarHandleOp's output. If it is anonymous, always
|
||||
// creates a new ID; otherwise, tries to reuse the existing ID for the
|
||||
// referenced variable if it exists, or creates a new one if not.
|
||||
int64_t GetOrCreateIdForVarHandle(TF::VarHandleOp handle, int64_t* next_id,
|
||||
llvm::StringMap<int64_t>* name_id_map) {
|
||||
// Always create a new ID for anonymous handle.
|
||||
if (IsResourceHandleAnonymous(handle)) return (*next_id)++;
|
||||
|
||||
auto name = GetVarHandleStringId(handle);
|
||||
auto emplace_res = name_id_map->try_emplace(name, *next_id);
|
||||
// New ID created, increment next_id.
|
||||
if (emplace_res.second) ++(*next_id);
|
||||
return emplace_res.first->second;
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
ResourceAliasAnalysis::ResourceAliasAnalysis(Operation* op) {
|
||||
auto func_op = llvm::dyn_cast<FuncOp>(op);
|
||||
if (!func_op) return;
|
||||
AnalyzeFunction(func_op);
|
||||
}
|
||||
|
||||
void ResourceAliasAnalysis::AnalyzeFunction(FuncOp func_op) {
|
||||
// This function populates resource_value_to_ids_.
|
||||
//
|
||||
// TODO(yuanzx): Pass variable aliasing information to functions so we can
|
||||
// properly resolve aliasing arguments.
|
||||
//
|
||||
// Before having that, we assume function arguments do not alias each other.
|
||||
int64_t next_unique_id = 0;
|
||||
for (auto arg : func_op.getArguments()) {
|
||||
if (!mlir::getElementTypeOrSelf(arg->getType()).isa<TF::ResourceType>())
|
||||
continue;
|
||||
resource_value_to_ids_[arg].insert(next_unique_id++);
|
||||
}
|
||||
llvm::StringMap<int64_t> var_handle_name_id_map;
|
||||
auto forward_input_to_output = [&](Value* operand, Value* result) {
|
||||
if (!mlir::getElementTypeOrSelf(result->getType()).isa<TF::ResourceType>())
|
||||
return;
|
||||
auto operand_it = resource_value_to_ids_.find(operand);
|
||||
assert(operand_it != resource_value_to_ids_.end() &&
|
||||
"A resource-type output does not have the corresponding "
|
||||
"resource-type input.");
|
||||
resource_value_to_ids_[result].insert(operand_it->getSecond().begin(),
|
||||
operand_it->getSecond().end());
|
||||
};
|
||||
// TODO(yuanzx): Consider control-flow ops.
|
||||
func_op.walk([&](Operation* op) {
|
||||
if (auto var_handle = llvm::dyn_cast<TF::VarHandleOp>(op)) {
|
||||
resource_value_to_ids_[var_handle.resource()].insert(
|
||||
GetOrCreateIdForVarHandle(var_handle, &next_unique_id,
|
||||
&var_handle_name_id_map));
|
||||
} else if (llvm::isa<TF::IdentityNOp>(op) ||
|
||||
llvm::isa<TF::IdentityOp>(op)) {
|
||||
for (auto operand_and_result :
|
||||
llvm::zip(op->getOperands(), op->getResults())) {
|
||||
forward_input_to_output(std::get<0>(operand_and_result),
|
||||
std::get<1>(operand_and_result));
|
||||
}
|
||||
} else {
|
||||
for (auto result : op->getResults()) {
|
||||
if (!mlir::getElementTypeOrSelf(result->getType())
|
||||
.isa<TF::ResourceType>())
|
||||
continue;
|
||||
resource_value_to_ids_[result].insert(kUnknownResourceId);
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
bool ResourceAliasAnalysis::IsUnknownResource(const Value* resource) const {
|
||||
auto it = resource_value_to_ids_.find(resource);
|
||||
assert(it != resource_value_to_ids_.end() && !it->getSecond().empty());
|
||||
// The set is sorted so we only need to check the first element since
|
||||
// kUnknownResourceId < 0.
|
||||
static_assert(kUnknownResourceId < 0,
|
||||
"kUnknownResourceId should be negative");
|
||||
return *it->getSecond().begin() == kUnknownResourceId;
|
||||
}
|
||||
|
||||
const llvm::SmallSet<int64_t, 8>& ResourceAliasAnalysis::GetResourceUniqueIds(
|
||||
const Value* resource) const {
|
||||
auto it = resource_value_to_ids_.find(resource);
|
||||
assert(it != resource_value_to_ids_.end() && "Unseen resource was queried");
|
||||
return it->getSecond();
|
||||
}
|
||||
|
||||
namespace {
|
||||
|
||||
// Returns a set that contains only kUnknownResourceId.
|
||||
llvm::SmallDenseSet<int64_t, 8> UnknownResourceSet() {
|
||||
llvm::SmallDenseSet<int64_t, 8> unknown_set;
|
||||
unknown_set.insert(kUnknownResourceId);
|
||||
return unknown_set;
|
||||
}
|
||||
|
||||
// Returns all resources that could be accessed by op, or UnknownResourceSet()
|
||||
// if we cannot find all of them.
|
||||
llvm::SmallDenseSet<int64_t, 8> FindAccessedResources(
|
||||
Operation* op, const ResourceAliasAnalysis& alias_analysis) {
|
||||
llvm::SmallDenseSet<int64_t, 8> resources;
|
||||
|
||||
for (auto operand : op->getOperands()) {
|
||||
if (!mlir::getElementTypeOrSelf(operand->getType()).isa<TF::ResourceType>())
|
||||
continue;
|
||||
if (alias_analysis.IsUnknownResource(operand)) return UnknownResourceSet();
|
||||
const auto& ids = alias_analysis.GetResourceUniqueIds(operand);
|
||||
resources.insert(ids.begin(), ids.end());
|
||||
}
|
||||
for (auto result : op->getResults()) {
|
||||
if (!mlir::getElementTypeOrSelf(result->getType()).isa<TF::ResourceType>())
|
||||
continue;
|
||||
if (alias_analysis.IsUnknownResource(result)) return UnknownResourceSet();
|
||||
const auto& ids = alias_analysis.GetResourceUniqueIds(result);
|
||||
resources.insert(ids.begin(), ids.end());
|
||||
}
|
||||
return resources;
|
||||
}
|
||||
|
||||
// Returns an XlaResourceOpInfo (or nullptr if it does not exist) that specifies
|
||||
// the resource access type of the op. It tells whether the op is read only,
|
||||
// etc.
|
||||
//
|
||||
// TODO(yuanzx): Define this information in a different place. Currently we use
|
||||
// tensorflow/compiler/tf2xla/resource_operation_table.h.
|
||||
const tensorflow::XlaResourceOpInfo* GetResourceInfoForOp(Operation* op) {
|
||||
auto op_name = op->getName().getStringRef().str();
|
||||
if (op->getName().getDialect() !=
|
||||
TF::TensorFlowDialect::getDialectNamespace()) {
|
||||
return nullptr;
|
||||
}
|
||||
return tensorflow::GetResourceOpInfoForOp(
|
||||
op->getName().getStringRef().split('.').second.str());
|
||||
}
|
||||
|
||||
// Returns whether `op` accesses resources and it is known to be read-only.
|
||||
bool OpIsReadOnly(Operation* op) {
|
||||
auto resource_op_info = GetResourceInfoForOp(op);
|
||||
return resource_op_info &&
|
||||
resource_op_info->kind() == tensorflow::XlaResourceOpKind::kRead;
|
||||
}
|
||||
|
||||
// Returns if `op` is a resource declaration.
|
||||
bool OpIsDeclaration(Operation* op,
|
||||
const ResourceAliasAnalysis& alias_analysis) {
|
||||
// TODO(yuanzx): Add other types of resources.
|
||||
return llvm::isa<TF::VarHandleOp>(op) ||
|
||||
((llvm::isa<TF::IdentityNOp>(op) || llvm::isa<TF::IdentityOp>(op)) &&
|
||||
!FindAccessedResources(op, alias_analysis).empty());
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
void SideEffectAnalysis::TrackAccess(int64_t resource_id, Operation* op,
|
||||
bool read_only) {
|
||||
if (resource_id == kUnknownResourceId) {
|
||||
if (read_only) {
|
||||
// New unknown read is not tracked by any known resource access.
|
||||
for (auto& entry : per_resource_access_info_) {
|
||||
entry.getSecond().tracked_last_unknown_read = false;
|
||||
}
|
||||
} else {
|
||||
// Unknown write can clear all other tracked information, since it acts
|
||||
// like a barrier.
|
||||
per_resource_access_info_.clear();
|
||||
}
|
||||
}
|
||||
auto& info = per_resource_access_info_[resource_id];
|
||||
if (read_only) {
|
||||
info.reads_since_last_write.push_back(op);
|
||||
// Resource read must have carried control dependencies of unknown write.
|
||||
info.tracked_last_unknown_write = true;
|
||||
} else {
|
||||
// Resource write must have carried control dependencies of unknown access.
|
||||
info.tracked_last_unknown_write = true;
|
||||
info.tracked_last_unknown_read = true;
|
||||
info.last_write = op;
|
||||
info.reads_since_last_write.clear();
|
||||
}
|
||||
}
|
||||
|
||||
void SideEffectAnalysis::AddPredecessorsForAccess(int64_t resource_id,
|
||||
Operation* op,
|
||||
bool read_only) {
|
||||
auto it = per_resource_access_info_.find(resource_id);
|
||||
if (it == per_resource_access_info_.end()) return;
|
||||
const auto& access_info = it->getSecond();
|
||||
auto& control_predecessors = control_predecessors_[op];
|
||||
bool read_tracked = false;
|
||||
if (!read_only) {
|
||||
control_predecessors.insert(access_info.reads_since_last_write.begin(),
|
||||
access_info.reads_since_last_write.end());
|
||||
read_tracked = !access_info.reads_since_last_write.empty();
|
||||
}
|
||||
if (access_info.last_write && !read_tracked) {
|
||||
control_predecessors.insert(access_info.last_write);
|
||||
}
|
||||
}
|
||||
|
||||
void SideEffectAnalysis::AnalyzeFunction(
|
||||
FuncOp func_op, const ResourceAliasAnalysis& alias_analysis) {
|
||||
// This function populates control_predecessors_ and control_successors_ by
|
||||
// walking through func_op's body, and tracking resource accesses in
|
||||
// per_resource_access_info_.
|
||||
|
||||
// Returns whether an access to `resource` can skip control edges from
|
||||
// prevoius accesses to unknown resources, due to that earlier accesses to
|
||||
// `resource` already indirectly tracked previous accesses to uknown
|
||||
// resources. `read_only` specifies the type of access of the current op being
|
||||
// considered.
|
||||
auto unknown_access_indirectly_tracked_by_resource = [&](int64_t resource,
|
||||
bool read_only) {
|
||||
auto it = per_resource_access_info_.find(resource);
|
||||
if (it == per_resource_access_info_.end()) return false;
|
||||
auto unknown_it = per_resource_access_info_.find(kUnknownResourceId);
|
||||
const bool no_unknown_read =
|
||||
unknown_it == per_resource_access_info_.end() ||
|
||||
unknown_it->getSecond().reads_since_last_write.empty();
|
||||
return read_only
|
||||
? it->second.tracked_last_unknown_write
|
||||
: it->second.tracked_last_unknown_write &&
|
||||
(it->second.tracked_last_unknown_read || no_unknown_read);
|
||||
};
|
||||
|
||||
func_op.walk([&](Operation* op) {
|
||||
// We do not need explicit control edges for declaration ops.
|
||||
if (OpIsDeclaration(op, alias_analysis)) return;
|
||||
|
||||
auto resource_op_info = GetResourceInfoForOp(op);
|
||||
if (!resource_op_info && op->hasNoSideEffect()) return;
|
||||
|
||||
llvm::SmallDenseSet<int64_t, 8> resources =
|
||||
resource_op_info ? FindAccessedResources(op, alias_analysis)
|
||||
: UnknownResourceSet();
|
||||
assert(!resources.empty());
|
||||
const bool is_unknown = resources.count(kUnknownResourceId) > 0;
|
||||
const bool read_only = OpIsReadOnly(op);
|
||||
bool indirectly_tracked_unknown_access = false;
|
||||
// First add edges from known resources.
|
||||
if (is_unknown) {
|
||||
for (auto& entry : per_resource_access_info_) {
|
||||
if (entry.getFirst() == kUnknownResourceId) continue;
|
||||
AddPredecessorsForAccess(entry.getFirst(), op, read_only);
|
||||
indirectly_tracked_unknown_access |=
|
||||
unknown_access_indirectly_tracked_by_resource(entry.getFirst(),
|
||||
read_only);
|
||||
}
|
||||
} else {
|
||||
for (int64_t resource : resources) {
|
||||
AddPredecessorsForAccess(resource, op, read_only);
|
||||
indirectly_tracked_unknown_access |=
|
||||
unknown_access_indirectly_tracked_by_resource(resource, read_only);
|
||||
// Update access info for known resources.
|
||||
TrackAccess(resource, op, read_only);
|
||||
}
|
||||
}
|
||||
// If not indirectly tracked, add edges from the unknown resource.
|
||||
if (!indirectly_tracked_unknown_access) {
|
||||
AddPredecessorsForAccess(kUnknownResourceId, op, read_only);
|
||||
}
|
||||
if (is_unknown) {
|
||||
// Update access info for unknown resource.
|
||||
TrackAccess(kUnknownResourceId, op, read_only);
|
||||
}
|
||||
});
|
||||
|
||||
// Populate control_successors_ based on control_predecessors_.
|
||||
for (auto& entry : control_predecessors_) {
|
||||
auto op = entry.getFirst();
|
||||
for (auto predecessor : entry.getSecond()) {
|
||||
control_successors_[predecessor].insert(op);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
llvm::SmallVector<Operation*, 8> SideEffectAnalysis::DirectControlPredecessors(
|
||||
Operation* op, llvm::function_ref<bool(Operation*)> filter) const {
|
||||
llvm::SmallVector<Operation*, 8> result;
|
||||
auto it = control_predecessors_.find(op);
|
||||
if (it == control_predecessors_.end()) return result;
|
||||
result.reserve(it->getSecond().size());
|
||||
for (auto predecessor : it->getSecond()) {
|
||||
if (!filter || filter(predecessor)) result.push_back(predecessor);
|
||||
}
|
||||
llvm::sort(result,
|
||||
[](Operation* a, Operation* b) { return a->isBeforeInBlock(b); });
|
||||
return result;
|
||||
}
|
||||
|
||||
llvm::SmallVector<Operation*, 8> SideEffectAnalysis::DirectControlSuccessors(
|
||||
Operation* op, llvm::function_ref<bool(Operation*)> filter) const {
|
||||
llvm::SmallVector<Operation*, 8> result;
|
||||
auto it = control_successors_.find(op);
|
||||
if (it == control_successors_.end()) return result;
|
||||
result.reserve(it->getSecond().size());
|
||||
for (auto successor : it->getSecond()) {
|
||||
if (!filter || filter(successor)) result.push_back(successor);
|
||||
}
|
||||
llvm::sort(result,
|
||||
[](Operation* a, Operation* b) { return a->isBeforeInBlock(b); });
|
||||
return result;
|
||||
}
|
||||
|
||||
SideEffectAnalysis::SideEffectAnalysis(Operation* op) {
|
||||
auto func_op = llvm::dyn_cast<FuncOp>(op);
|
||||
if (!func_op) return;
|
||||
ResourceAliasAnalysis alias_analysis(op);
|
||||
AnalyzeFunction(func_op, alias_analysis);
|
||||
}
|
||||
|
||||
} // namespace TF
|
||||
} // namespace mlir
|
||||
@ -0,0 +1,125 @@
|
||||
/* Copyright 2019 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_MLIR_TENSORFLOW_ANALYSIS_SIDE_EFFECT_ANALYSIS_H_
|
||||
#define TENSORFLOW_COMPILER_MLIR_TENSORFLOW_ANALYSIS_SIDE_EFFECT_ANALYSIS_H_
|
||||
|
||||
#include <cstdint>
|
||||
#include <memory>
|
||||
|
||||
#include "llvm/ADT/SmallSet.h"
|
||||
#include "llvm/ADT/SmallVector.h"
|
||||
#include "llvm/ADT/StringMap.h"
|
||||
#include "mlir/IR/Function.h" // TF:local_config_mlir
|
||||
#include "mlir/IR/Operation.h" // TF:local_config_mlir
|
||||
#include "mlir/IR/Region.h" // TF:local_config_mlir
|
||||
#include "mlir/Support/LogicalResult.h" // TF:local_config_mlir
|
||||
|
||||
namespace mlir {
|
||||
namespace TF {
|
||||
|
||||
// An analysis that runs on a function and maps each resource-type value to a
|
||||
// set of unique int64_t IDs representing the possible resources it could alias.
|
||||
class ResourceAliasAnalysis {
|
||||
public:
|
||||
explicit ResourceAliasAnalysis(Operation* op);
|
||||
~ResourceAliasAnalysis() = default;
|
||||
ResourceAliasAnalysis(ResourceAliasAnalysis&&) = default;
|
||||
|
||||
// Returns if the analysis fails to resolve a resource-type value.
|
||||
bool IsUnknownResource(const Value* resource) const;
|
||||
|
||||
// Returns the set unique IDs which `resource` could alias. Requires that
|
||||
// IsUnknownResource(resource) == true.
|
||||
const llvm::SmallSet<int64_t, 8>& GetResourceUniqueIds(
|
||||
const Value* resource) const;
|
||||
|
||||
private:
|
||||
ResourceAliasAnalysis() = default;
|
||||
|
||||
// Runs the analysis on `func_op` and populates resource_value_to_ids_.
|
||||
void AnalyzeFunction(FuncOp func_op);
|
||||
|
||||
// Maps each resource-type value to a set of unique IDs that it could alias.
|
||||
llvm::SmallDenseMap<const Value*, llvm::SmallSet<int64_t, 8>, 8>
|
||||
resource_value_to_ids_;
|
||||
};
|
||||
|
||||
// An analysis that runs on a function and infers the control predecessors and
|
||||
// successors for each op, based on side-effects on known and unknown resources.
|
||||
// Side-effecting ops on uknown resources are conservatively treated as
|
||||
// interfering with all known resource op accesses. It distinguishes accesses
|
||||
// based on whether they are read-only, and read-only ops do not interfer with
|
||||
// each other.
|
||||
class SideEffectAnalysis {
|
||||
public:
|
||||
explicit SideEffectAnalysis(Operation* op);
|
||||
SideEffectAnalysis(SideEffectAnalysis&& other) = default;
|
||||
~SideEffectAnalysis() = default;
|
||||
|
||||
// Returns a vector of ops that are direct control predecessors of `op`,
|
||||
// sorted in program order. If `filter` is provided, only predecessors that
|
||||
// pass the filter (returning true) will be included.
|
||||
llvm::SmallVector<Operation*, 8> DirectControlPredecessors(
|
||||
Operation* op,
|
||||
llvm::function_ref<bool(Operation*)> filter = nullptr) const;
|
||||
|
||||
// Returns a vector of ops that are direct control successors of `op`, sorted
|
||||
// in program order. If `filter` is provided, only successors that pass the
|
||||
// filter (returning true) will be included.
|
||||
llvm::SmallVector<Operation*, 8> DirectControlSuccessors(
|
||||
Operation* op,
|
||||
llvm::function_ref<bool(Operation*)> filter = nullptr) const;
|
||||
|
||||
private:
|
||||
// Runs the analysis on `func_op` and populates control_predecessors_ and
|
||||
// control_successors_.
|
||||
void AnalyzeFunction(FuncOp func_op,
|
||||
const ResourceAliasAnalysis& alias_analysis);
|
||||
|
||||
// Updates control_predecessors_ for `op` that is being visted, on the given
|
||||
// `resource_id`.
|
||||
void AddPredecessorsForAccess(int64_t resource_id, Operation* op,
|
||||
bool read_only);
|
||||
|
||||
// Adds op's access to per_resource_access_info_.
|
||||
void TrackAccess(int64_t resource_id, Operation* op, bool read_only);
|
||||
|
||||
// Maps from an op to its control predecessors.
|
||||
llvm::SmallDenseMap<Operation*, llvm::SmallPtrSet<Operation*, 8>, 8>
|
||||
control_predecessors_;
|
||||
// Maps from an op to its control successors.
|
||||
llvm::SmallDenseMap<Operation*, llvm::SmallPtrSet<Operation*, 8>, 8>
|
||||
control_successors_;
|
||||
|
||||
// Internal per-resource data structure when we build the dependencies.
|
||||
struct PerResourceAcessInfo {
|
||||
// Last op that writes the resource before the current op being analyzed.
|
||||
Operation* last_write = nullptr;
|
||||
// Read ops since last_write before the current op being analyzed.
|
||||
llvm::SmallVector<Operation*, 8> reads_since_last_write;
|
||||
// Whether previous accesses of this resource already tracked last unknown
|
||||
// read/write.
|
||||
bool tracked_last_unknown_read = false;
|
||||
bool tracked_last_unknown_write = false;
|
||||
};
|
||||
llvm::SmallDenseMap<int64_t, PerResourceAcessInfo, 8>
|
||||
per_resource_access_info_;
|
||||
};
|
||||
|
||||
} // namespace TF
|
||||
} // namespace mlir
|
||||
|
||||
#endif // TENSORFLOW_COMPILER_MLIR_TENSORFLOW_ANALYSIS_SIDE_EFFECT_ANALYSIS_H_
|
||||
@ -0,0 +1,237 @@
|
||||
// RUN: tf-opt -split-input-file -tf-test-side-effect-analysis -verify-diagnostics %s | FileCheck %s --dump-input=fail
|
||||
|
||||
// Tests that the pass tracks control dependencies for reads/writes on the same
|
||||
// resource.
|
||||
|
||||
// CHECK-LABEL: func @non_aliasing_reads_writes
|
||||
func @non_aliasing_reads_writes(
|
||||
// expected-remark@above {{ID: 13}}
|
||||
// expected-remark@above {{Predecessors: {12}}}
|
||||
%arg0: tensor<*x!tf.resource<tensor<32xf32>>>,
|
||||
%arg1: tensor<*x!tf.resource<tensor<32xf32>>>,
|
||||
%arg2: tensor<32xf32>) -> (tensor<32xf32>) {
|
||||
%graph = tf_executor.graph {
|
||||
// expected-remark@above {{ID: 11}}
|
||||
// expected-remark@above {{Predecessors: {10}}}
|
||||
// expected-remark@above {{Successors: {12}}}
|
||||
// CHECK: tf_executor.island
|
||||
%island:2 = tf_executor.island {
|
||||
// expected-remark@above {{ID: 9}}
|
||||
// expected-remark@above {{Predecessors: {8}}}
|
||||
// expected-remark@above {{Successors: {10}}}
|
||||
%read0 = "tf.ReadVariableOp"(%arg0) : (tensor<*x!tf.resource<tensor<32xf32>>>) -> tensor<32xf32>
|
||||
// expected-remark@above {{ID: 0}}
|
||||
// expected-remark@above {{Successors: {1}}}
|
||||
"tf.AssignVariableOp"(%arg0, %arg2) : (tensor<*x!tf.resource<tensor<32xf32>>>, tensor<32xf32>) -> ()
|
||||
// expected-remark@above {{ID: 1}}
|
||||
// expected-remark@above {{Predecessors: {0}}}
|
||||
// expected-remark@above {{Successors: {6}}}
|
||||
%read1 = "tf.ReadVariableOp"(%arg1) : (tensor<*x!tf.resource<tensor<32xf32>>>) -> tensor<32xf32>
|
||||
// expected-remark@above {{ID: 2}}
|
||||
// expected-remark@above {{Successors: {5}}}
|
||||
%var_handle = "tf.VarHandleOp"() {container = "c", shared_name = "v0"} : () -> tensor<*x!tf.resource<tensor<32xf32>>>
|
||||
// expected-remark@above {{ID: 3}}
|
||||
%read2 = "tf.ReadVariableOp"(%var_handle) : (tensor<*x!tf.resource<tensor<32xf32>>>) -> tensor<32xf32>
|
||||
// expected-remark@above {{ID: 4}}
|
||||
// expected-remark@above {{Successors: {8}}}
|
||||
"tf.AssignVariableOp"(%arg1, %read0) : (tensor<*x!tf.resource<tensor<32xf32>>>, tensor<32xf32>) -> ()
|
||||
// expected-remark@above {{ID: 5}}
|
||||
// expected-remark@above {{Predecessors: {2}}}
|
||||
// expected-remark@above {{Successors: {8}}}
|
||||
"tf.AssignVariableOp"(%arg0, %read2) : (tensor<*x!tf.resource<tensor<32xf32>>>, tensor<32xf32>) -> ()
|
||||
// expected-remark@above {{ID: 6}}
|
||||
// expected-remark@above {{Predecessors: {1}}}
|
||||
// expected-remark@above {{Successors: {7}}}
|
||||
%read3 = "tf.ReadVariableOp"(%arg0) : (tensor<*x!tf.resource<tensor<32xf32>>>) -> tensor<32xf32>
|
||||
// expected-remark@above {{ID: 7}}
|
||||
// expected-remark@above {{Predecessors: {6}}}
|
||||
// expected-remark@above {{Successors: {8}}}
|
||||
tf_executor.yield %read3 : tensor<32xf32>
|
||||
// expected-remark@above {{ID: 8}}
|
||||
// expected-remark@above {{Predecessors: {4,5,7}}}
|
||||
// expected-remark@above {{Successors: {9}}}
|
||||
}
|
||||
tf_executor.fetch %island#0 : tensor<32xf32>
|
||||
// expected-remark@above {{ID: 10}}
|
||||
// expected-remark@above {{Predecessors: {9}}}
|
||||
// expected-remark@above {{Successors: {11}}}
|
||||
}
|
||||
return %graph : tensor<32xf32>
|
||||
// expected-remark@above {{ID: 12}}
|
||||
// expected-remark@above {{Predecessors: {11}}}
|
||||
// expected-remark@above {{Successors: {13}}}
|
||||
}
|
||||
|
||||
// -----
|
||||
|
||||
// Tests that the pass tracks control dependencies for reads/writes on the two
|
||||
// resource handles that refer to the same variable.
|
||||
|
||||
// CHECK-LABEL: func @aliasing_reads_writes
|
||||
func @aliasing_reads_writes(%arg0: tensor<32xf32>) -> () {
|
||||
// expected-remark@above {{ID: 14}}
|
||||
// expected-remark@above {{Predecessors: {13}}}
|
||||
tf_executor.graph {
|
||||
// expected-remark@above {{ID: 12}}
|
||||
// expected-remark@above {{Predecessors: {11}}}
|
||||
// expected-remark@above {{Successors: {13}}}
|
||||
// CHECK: tf_executor.island
|
||||
%island = tf_executor.island {
|
||||
// expected-remark@above {{ID: 10}}
|
||||
// expected-remark@above {{Predecessors: {9}}}
|
||||
// expected-remark@above {{Successors: {11}}}
|
||||
%vh0 = "tf.VarHandleOp"() {container = "c", shared_name = "v0"} : () -> tensor<*x!tf.resource<tensor<32xf32>>>
|
||||
// expected-remark@above {{ID: 0}}
|
||||
%vh1 = "tf.VarHandleOp"() {container = "c", shared_name = "v0"} : () -> tensor<*x!tf.resource<tensor<32xf32>>>
|
||||
// expected-remark@above {{ID: 1}}
|
||||
%vh1_id:2 = "tf.IdentityN"(%vh1, %arg0) : (tensor<*x!tf.resource<tensor<32xf32>>>, tensor<32xf32>) -> (tensor<*x!tf.resource<tensor<32xf32>>>, tensor<32xf32>)
|
||||
// expected-remark@above {{ID: 2}}
|
||||
%read0 = "tf.ReadVariableOp"(%vh0) : (tensor<*x!tf.resource<tensor<32xf32>>>) -> tensor<32xf32>
|
||||
// expected-remark@above {{ID: 3}}
|
||||
// expected-remark@above {{Successors: {4}}}
|
||||
"tf.AssignVariableOp"(%vh1_id#0, %arg0) : (tensor<*x!tf.resource<tensor<32xf32>>>, tensor<32xf32>) -> ()
|
||||
// expected-remark@above {{ID: 4}}
|
||||
// expected-remark@above {{Predecessors: {3}}}
|
||||
// expected-remark@above {{Successors: {5,6}}}
|
||||
%read1 = "tf.ReadVariableOp"(%vh0) : (tensor<*x!tf.resource<tensor<32xf32>>>) -> tensor<32xf32>
|
||||
// expected-remark@above {{ID: 5}}
|
||||
// expected-remark@above {{Predecessors: {4}}}
|
||||
// expected-remark@above {{Successors: {7}}}
|
||||
%read2 = "tf.ReadVariableOp"(%vh1) : (tensor<*x!tf.resource<tensor<32xf32>>>) -> tensor<32xf32>
|
||||
// expected-remark@above {{ID: 6}}
|
||||
// expected-remark@above {{Predecessors: {4}}}
|
||||
// expected-remark@above {{Successors: {7}}}
|
||||
"tf.AssignVariableOp"(%vh0, %read2) : (tensor<*x!tf.resource<tensor<32xf32>>>, tensor<32xf32>) -> ()
|
||||
// expected-remark@above {{ID: 7}}
|
||||
// expected-remark@above {{Predecessors: {5,6}}}
|
||||
// expected-remark@above {{Successors: {8}}}
|
||||
"tf.AssignVariableOp"(%vh1_id#0, %read1) : (tensor<*x!tf.resource<tensor<32xf32>>>, tensor<32xf32>) -> ()
|
||||
// expected-remark@above {{ID: 8}}
|
||||
// expected-remark@above {{Predecessors: {7}}}
|
||||
// expected-remark@above {{Successors: {9}}}
|
||||
tf_executor.yield
|
||||
// expected-remark@above {{ID: 9}}
|
||||
// expected-remark@above {{Predecessors: {8}}}
|
||||
// expected-remark@above {{Successors: {10}}}
|
||||
}
|
||||
tf_executor.fetch %island : !tf_executor.control
|
||||
// expected-remark@above {{ID: 11}}
|
||||
// expected-remark@above {{Predecessors: {10}}}
|
||||
// expected-remark@above {{Successors: {12}}}
|
||||
}
|
||||
return
|
||||
// expected-remark@above {{ID: 13}}
|
||||
// expected-remark@above {{Predecessors: {12}}}
|
||||
// expected-remark@above {{Successors: {14}}}
|
||||
}
|
||||
|
||||
// -----
|
||||
|
||||
// Tests that the pass tracks control dependencies for side-effecting on unknown
|
||||
// resources.
|
||||
|
||||
// CHECK-LABEL: func @unknown_side_effecting_op
|
||||
func @unknown_side_effecting_op(%arg0: tensor<32xf32>) -> () {
|
||||
// expected-remark@above {{ID: 13}}
|
||||
// expected-remark@above {{Predecessors: {12}}}
|
||||
tf_executor.graph {
|
||||
// expected-remark@above {{ID: 11}}
|
||||
// expected-remark@above {{Predecessors: {10}}}
|
||||
// expected-remark@above {{Successors: {12}}}
|
||||
// CHECK: tf_executor.island
|
||||
%island = tf_executor.island {
|
||||
// expected-remark@above {{ID: 9}}
|
||||
// expected-remark@above {{Predecessors: {8}}}
|
||||
// expected-remark@above {{Successors: {10}}}
|
||||
%vh0 = "tf.VarHandleOp"() {container = "c", shared_name = "v0"} : () -> tensor<*x!tf.resource<tensor<32xf32>>>
|
||||
// expected-remark@above {{ID: 0}}
|
||||
%vh1 = "tf.VarHandleOp"() {container = "c", shared_name = "v1"} : () -> tensor<*x!tf.resource<tensor<32xf32>>>
|
||||
// expected-remark@above {{ID: 1}}
|
||||
%read0 = "tf.ReadVariableOp"(%vh0) : (tensor<*x!tf.resource<tensor<32xf32>>>) -> tensor<32xf32>
|
||||
// expected-remark@above {{ID: 2}}
|
||||
// expected-remark@above {{Successors: {4}}}
|
||||
"tf.AssignVariableOp"(%vh1, %arg0) : (tensor<*x!tf.resource<tensor<32xf32>>>, tensor<32xf32>) -> ()
|
||||
// expected-remark@above {{ID: 3}}
|
||||
// expected-remark@above {{Successors: {4}}}
|
||||
"tf._UnknownSideEffectingOp_"() : () -> ()
|
||||
// expected-remark@above {{ID: 4}}
|
||||
// expected-remark@above {{Predecessors: {2,3}}}
|
||||
// expected-remark@above {{Successors: {5,6}}}
|
||||
%read1 = "tf.ReadVariableOp"(%vh1) : (tensor<*x!tf.resource<tensor<32xf32>>>) -> tensor<32xf32>
|
||||
// expected-remark@above {{ID: 5}}
|
||||
// expected-remark@above {{Predecessors: {4}}}
|
||||
// expected-remark@above {{Successors: {7}}}
|
||||
"tf.AssignVariableOp"(%vh0, %read1) : (tensor<*x!tf.resource<tensor<32xf32>>>, tensor<32xf32>) -> ()
|
||||
// expected-remark@above {{ID: 6}}
|
||||
// expected-remark@above {{Predecessors: {4}}}
|
||||
// expected-remark@above {{Successors: {8}}}
|
||||
"tf.AssignVariableOp"(%vh1, %read0) : (tensor<*x!tf.resource<tensor<32xf32>>>, tensor<32xf32>) -> ()
|
||||
// expected-remark@above {{ID: 7}}
|
||||
// expected-remark@above {{Predecessors: {5}}}
|
||||
// expected-remark@above {{Successors: {8}}}
|
||||
tf_executor.yield
|
||||
// expected-remark@above {{ID: 8}}
|
||||
// expected-remark@above {{Predecessors: {6,7}}}
|
||||
// expected-remark@above {{Successors: {9}}}
|
||||
}
|
||||
tf_executor.fetch %island : !tf_executor.control
|
||||
// expected-remark@above {{ID: 10}}
|
||||
// expected-remark@above {{Predecessors: {9}}}
|
||||
// expected-remark@above {{Successors: {11}}}
|
||||
}
|
||||
return
|
||||
// expected-remark@above {{ID: 12}}
|
||||
// expected-remark@above {{Predecessors: {11}}}
|
||||
// expected-remark@above {{Successors: {13}}}
|
||||
}
|
||||
|
||||
// -----
|
||||
|
||||
// Tests that the pass tracks control dependencies for read-only ops on unknown
|
||||
// resources.
|
||||
|
||||
// CHECK-LABEL: func @read_only_unknown_resource
|
||||
func @read_only_unknown_resource(%arg0: tensor<32xf32>) -> () {
|
||||
// expected-remark@above {{ID: 10}}
|
||||
// expected-remark@above {{Predecessors: {9}}}
|
||||
tf_executor.graph {
|
||||
// expected-remark@above {{ID: 8}}
|
||||
// expected-remark@above {{Predecessors: {7}}}
|
||||
// expected-remark@above {{Successors: {9}}}
|
||||
// CHECK: tf_executor.island
|
||||
%island = tf_executor.island {
|
||||
// expected-remark@above {{ID: 6}}
|
||||
// expected-remark@above {{Predecessors: {5}}}
|
||||
// expected-remark@above {{Successors: {7}}}
|
||||
%vh0 = "tf._UnknownSideEffectingOp_"() : () -> tensor<*x!tf.resource<tensor<32xf32>>>
|
||||
// expected-remark@above {{ID: 0}}
|
||||
// expected-remark@above {{Successors: {2,3}}}
|
||||
%vh1 = "tf.VarHandleOp"() {container = "c", shared_name = "v1"} : () -> tensor<*x!tf.resource<tensor<32xf32>>>
|
||||
// expected-remark@above {{ID: 1}}
|
||||
%read0 = "tf.ReadVariableOp"(%vh0) : (tensor<*x!tf.resource<tensor<32xf32>>>) -> tensor<32xf32>
|
||||
// expected-remark@above {{ID: 2}}
|
||||
// expected-remark@above {{Predecessors: {0}}}
|
||||
// expected-remark@above {{Successors: {4}}}
|
||||
%read1 = "tf.ReadVariableOp"(%vh1) : (tensor<*x!tf.resource<tensor<32xf32>>>) -> tensor<32xf32>
|
||||
// expected-remark@above {{ID: 3}}
|
||||
// expected-remark@above {{Predecessors: {0}}}
|
||||
// expected-remark@above {{Successors: {4}}}
|
||||
"tf.AssignVariableOp"(%vh1, %read0) : (tensor<*x!tf.resource<tensor<32xf32>>>, tensor<32xf32>) -> ()
|
||||
// expected-remark@above {{ID: 4}}
|
||||
// expected-remark@above {{Predecessors: {2,3}}}
|
||||
// expected-remark@above {{Successors: {5}}}
|
||||
tf_executor.yield
|
||||
// expected-remark@above {{ID: 5}}
|
||||
// expected-remark@above {{Predecessors: {4}}}
|
||||
// expected-remark@above {{Successors: {6}}}
|
||||
}
|
||||
tf_executor.fetch %island : !tf_executor.control
|
||||
// expected-remark@above {{ID: 7}}
|
||||
// expected-remark@above {{Predecessors: {6}}}
|
||||
// expected-remark@above {{Successors: {8}}}
|
||||
}
|
||||
return
|
||||
// expected-remark@above {{ID: 9}}
|
||||
// expected-remark@above {{Predecessors: {8}}}
|
||||
// expected-remark@above {{Successors: {10}}}
|
||||
}
|
||||
@ -0,0 +1,77 @@
|
||||
/* Copyright 2019 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 <cstddef>
|
||||
#include <cstdint>
|
||||
#include <string>
|
||||
#include <utility>
|
||||
|
||||
#include "llvm/ADT/DenseMap.h"
|
||||
#include "llvm/ADT/STLExtras.h"
|
||||
#include "llvm/ADT/SmallVector.h"
|
||||
#include "llvm/Support/Debug.h"
|
||||
#include "mlir/Pass/Pass.h" // TF:local_config_mlir
|
||||
#include "mlir/Pass/PassManager.h" // TF:local_config_mlir
|
||||
#include "mlir/Support/LLVM.h" // TF:local_config_mlir
|
||||
#include "mlir/Transforms/Passes.h" // TF:local_config_mlir
|
||||
#include "mlir/Transforms/RegionUtils.h" // TF:local_config_mlir
|
||||
#include "tensorflow/compiler/mlir/tensorflow/analysis/side_effect_analysis.h"
|
||||
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
|
||||
#include "tensorflow/compiler/mlir/tensorflow/utils/error_util.h"
|
||||
|
||||
namespace mlir {
|
||||
namespace tf_executor {
|
||||
|
||||
namespace {
|
||||
|
||||
// A pass that adds "Predecessors" and "Successors" remarks for each op based on
|
||||
// SideEffectAnalysis result. For testing purpose only.
|
||||
struct TestSideEffectAnalysis
|
||||
: public mlir::FunctionPass<TestSideEffectAnalysis> {
|
||||
void runOnFunction() override {
|
||||
int64_t next_id = 0;
|
||||
llvm::SmallDenseMap<Operation*, int64_t, 8> ids;
|
||||
getFunction().walk([&](Operation* op) {
|
||||
ids[op] = next_id++;
|
||||
op->emitRemark("ID: ") << ids[op];
|
||||
});
|
||||
auto join_ids = [&](const llvm::ArrayRef<Operation*> ops) {
|
||||
llvm::SmallVector<std::string, 8> id_vec;
|
||||
id_vec.reserve(ops.size());
|
||||
for (auto op : ops) id_vec.push_back(std::to_string(ids[op]));
|
||||
return llvm::join(id_vec, ",");
|
||||
};
|
||||
auto& analysis = getAnalysis<TF::SideEffectAnalysis>();
|
||||
getFunction().walk([&](Operation* op) {
|
||||
if (!analysis.DirectControlPredecessors(op).empty()) {
|
||||
op->emitRemark("Predecessors: ")
|
||||
<< "{" << join_ids(analysis.DirectControlPredecessors(op)) << "}";
|
||||
}
|
||||
if (!analysis.DirectControlSuccessors(op).empty()) {
|
||||
op->emitRemark("Successors: ")
|
||||
<< "{" << join_ids(analysis.DirectControlSuccessors(op)) << "}";
|
||||
}
|
||||
});
|
||||
}
|
||||
};
|
||||
|
||||
static mlir::PassRegistration<TestSideEffectAnalysis> pass(
|
||||
"tf-test-side-effect-analysis",
|
||||
"Add remarks based on side-effect analysis result, for testing purpose.");
|
||||
|
||||
} // anonymous namespace
|
||||
|
||||
} // namespace tf_executor
|
||||
} // namespace mlir
|
||||
Loading…
x
Reference in New Issue
Block a user