From fe77f7925607fbbeae9bf1a40335aeb7ef71ba83 Mon Sep 17 00:00:00 2001 From: Yuanzhong Xu Date: Wed, 27 Nov 2019 12:29:01 -0800 Subject: [PATCH] [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 --- tensorflow/compiler/mlir/tensorflow/BUILD | 18 + .../analysis/side_effect_analysis.cc | 374 ++++++++++++++++++ .../analysis/side_effect_analysis.h | 125 ++++++ .../tests/side-effect-analysis-test.mlir | 237 +++++++++++ .../transforms/test_side_effect_analysis.cc | 77 ++++ 5 files changed, 831 insertions(+) create mode 100644 tensorflow/compiler/mlir/tensorflow/analysis/side_effect_analysis.cc create mode 100644 tensorflow/compiler/mlir/tensorflow/analysis/side_effect_analysis.h create mode 100644 tensorflow/compiler/mlir/tensorflow/tests/side-effect-analysis-test.mlir create mode 100644 tensorflow/compiler/mlir/tensorflow/transforms/test_side_effect_analysis.cc diff --git a/tensorflow/compiler/mlir/tensorflow/BUILD b/tensorflow/compiler/mlir/tensorflow/BUILD index 683c66eced3..7cfa802b1c3 100644 --- a/tensorflow/compiler/mlir/tensorflow/BUILD +++ b/tensorflow/compiler/mlir/tensorflow/BUILD @@ -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", + ], +) diff --git a/tensorflow/compiler/mlir/tensorflow/analysis/side_effect_analysis.cc b/tensorflow/compiler/mlir/tensorflow/analysis/side_effect_analysis.cc new file mode 100644 index 00000000000..8d43c9330d0 --- /dev/null +++ b/tensorflow/compiler/mlir/tensorflow/analysis/side_effect_analysis.cc @@ -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 +#include + +#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("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* 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(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()) + continue; + resource_value_to_ids_[arg].insert(next_unique_id++); + } + llvm::StringMap var_handle_name_id_map; + auto forward_input_to_output = [&](Value* operand, Value* result) { + if (!mlir::getElementTypeOrSelf(result->getType()).isa()) + 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(op)) { + resource_value_to_ids_[var_handle.resource()].insert( + GetOrCreateIdForVarHandle(var_handle, &next_unique_id, + &var_handle_name_id_map)); + } else if (llvm::isa(op) || + llvm::isa(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()) + 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& 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 UnknownResourceSet() { + llvm::SmallDenseSet 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 FindAccessedResources( + Operation* op, const ResourceAliasAnalysis& alias_analysis) { + llvm::SmallDenseSet resources; + + for (auto operand : op->getOperands()) { + if (!mlir::getElementTypeOrSelf(operand->getType()).isa()) + 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()) + 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(op) || + ((llvm::isa(op) || llvm::isa(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 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 SideEffectAnalysis::DirectControlPredecessors( + Operation* op, llvm::function_ref filter) const { + llvm::SmallVector 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 SideEffectAnalysis::DirectControlSuccessors( + Operation* op, llvm::function_ref filter) const { + llvm::SmallVector 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(op); + if (!func_op) return; + ResourceAliasAnalysis alias_analysis(op); + AnalyzeFunction(func_op, alias_analysis); +} + +} // namespace TF +} // namespace mlir diff --git a/tensorflow/compiler/mlir/tensorflow/analysis/side_effect_analysis.h b/tensorflow/compiler/mlir/tensorflow/analysis/side_effect_analysis.h new file mode 100644 index 00000000000..5eee28a6ae0 --- /dev/null +++ b/tensorflow/compiler/mlir/tensorflow/analysis/side_effect_analysis.h @@ -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 +#include + +#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& 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, 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 DirectControlPredecessors( + Operation* op, + llvm::function_ref 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 DirectControlSuccessors( + Operation* op, + llvm::function_ref 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, 8> + control_predecessors_; + // Maps from an op to its control successors. + llvm::SmallDenseMap, 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 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 + per_resource_access_info_; +}; + +} // namespace TF +} // namespace mlir + +#endif // TENSORFLOW_COMPILER_MLIR_TENSORFLOW_ANALYSIS_SIDE_EFFECT_ANALYSIS_H_ diff --git a/tensorflow/compiler/mlir/tensorflow/tests/side-effect-analysis-test.mlir b/tensorflow/compiler/mlir/tensorflow/tests/side-effect-analysis-test.mlir new file mode 100644 index 00000000000..c6eb4663e57 --- /dev/null +++ b/tensorflow/compiler/mlir/tensorflow/tests/side-effect-analysis-test.mlir @@ -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>>, + %arg1: tensor<*x!tf.resource>>, + %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> + // expected-remark@above {{ID: 0}} + // expected-remark@above {{Successors: {1}}} + "tf.AssignVariableOp"(%arg0, %arg2) : (tensor<*x!tf.resource>>, 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> + // expected-remark@above {{ID: 2}} + // expected-remark@above {{Successors: {5}}} + %var_handle = "tf.VarHandleOp"() {container = "c", shared_name = "v0"} : () -> tensor<*x!tf.resource>> + // expected-remark@above {{ID: 3}} + %read2 = "tf.ReadVariableOp"(%var_handle) : (tensor<*x!tf.resource>>) -> tensor<32xf32> + // expected-remark@above {{ID: 4}} + // expected-remark@above {{Successors: {8}}} + "tf.AssignVariableOp"(%arg1, %read0) : (tensor<*x!tf.resource>>, 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>) -> () + // 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> + // 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>> + // expected-remark@above {{ID: 0}} + %vh1 = "tf.VarHandleOp"() {container = "c", shared_name = "v0"} : () -> tensor<*x!tf.resource>> + // expected-remark@above {{ID: 1}} + %vh1_id:2 = "tf.IdentityN"(%vh1, %arg0) : (tensor<*x!tf.resource>>, tensor<32xf32>) -> (tensor<*x!tf.resource>>, tensor<32xf32>) + // expected-remark@above {{ID: 2}} + %read0 = "tf.ReadVariableOp"(%vh0) : (tensor<*x!tf.resource>>) -> tensor<32xf32> + // expected-remark@above {{ID: 3}} + // expected-remark@above {{Successors: {4}}} + "tf.AssignVariableOp"(%vh1_id#0, %arg0) : (tensor<*x!tf.resource>>, 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> + // 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> + // 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>) -> () + // 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>) -> () + // 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>> + // expected-remark@above {{ID: 0}} + %vh1 = "tf.VarHandleOp"() {container = "c", shared_name = "v1"} : () -> tensor<*x!tf.resource>> + // expected-remark@above {{ID: 1}} + %read0 = "tf.ReadVariableOp"(%vh0) : (tensor<*x!tf.resource>>) -> tensor<32xf32> + // expected-remark@above {{ID: 2}} + // expected-remark@above {{Successors: {4}}} + "tf.AssignVariableOp"(%vh1, %arg0) : (tensor<*x!tf.resource>>, 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> + // 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>) -> () + // 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>) -> () + // 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>> + // expected-remark@above {{ID: 0}} + // expected-remark@above {{Successors: {2,3}}} + %vh1 = "tf.VarHandleOp"() {container = "c", shared_name = "v1"} : () -> tensor<*x!tf.resource>> + // expected-remark@above {{ID: 1}} + %read0 = "tf.ReadVariableOp"(%vh0) : (tensor<*x!tf.resource>>) -> 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> + // 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>) -> () + // 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}}} +} diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/test_side_effect_analysis.cc b/tensorflow/compiler/mlir/tensorflow/transforms/test_side_effect_analysis.cc new file mode 100644 index 00000000000..f0b7964389d --- /dev/null +++ b/tensorflow/compiler/mlir/tensorflow/transforms/test_side_effect_analysis.cc @@ -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 +#include +#include +#include + +#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 { + void runOnFunction() override { + int64_t next_id = 0; + llvm::SmallDenseMap ids; + getFunction().walk([&](Operation* op) { + ids[op] = next_id++; + op->emitRemark("ID: ") << ids[op]; + }); + auto join_ids = [&](const llvm::ArrayRef ops) { + llvm::SmallVector 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(); + 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 pass( + "tf-test-side-effect-analysis", + "Add remarks based on side-effect analysis result, for testing purpose."); + +} // anonymous namespace + +} // namespace tf_executor +} // namespace mlir