experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Add a new pass for promoting VarHandle ops to TF saved model arguments

Browse Source 
PiperOrigin-RevId: 314415654
Change-Id: I0dccc4dd9c3c20625950816d341ce910384d3906
This commit is contained in:
Jaesung Chung 2020-06-02 15:25:48 -07:00 committed by TensorFlower Gardener
parent e410915945
commit 96dfb6f4a8
3 changed files with 37 additions and 24 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
tensorflow/compiler/mlir/tensorflow/tests/promote_var_handles_to_args.mlir
View File

@ -1,4 +1,5 @@
// RUN: tf-opt %s -split-input-file -verify-diagnostics -tf-promote-var-handles-to-args | FileCheck %s -dump-input-on-failure
// Run a pass for promoting tf.VarHandleOps to function arguments in a format of TensorFlowSavedModelDialect.
// RUN: tf-opt %s -split-input-file -verify-diagnostics -tf-saved-model-promote-var-handles-to-args | FileCheck %s -dump-input-on-failure
// Tests main function with multiple blocks.
@ -11,24 +12,27 @@ func @main() {
// -----
"tf_saved_model.global_tensor"() {sym_name = "x", type = tensor<f32>, value = dense<1.67482901> : tensor<f32>} : () -> ()
"tf_saved_model.global_tensor"() {sym_name = "y", type = tensor<i32>, value = dense<0> : tensor<i32>} : () -> ()
// CHECK-LABEL: func @no_args
// CHECK-SAME: (%arg0: tensor<!tf.resource> {tf.resource_name = "x"})
// CHECK-SAME: (%arg0: tensor<!tf.resource<tensor<f32>>> {tf_saved_model.bound_input = @x})
// CHECK-NOT: "tf.VarHandleOp"
func @no_args() {
%0 = "tf.VarHandleOp"() {container = "", shape = "tfshape$", shared_name = "x"} : () -> tensor<!tf.resource>
%0 = "tf.VarHandleOp"() {container = "", shape = "tfshape$", shared_name = "x"} : () -> tensor<!tf.resource<tensor<f32>>>
return
}
// CHECK-LABEL: func @some_args
// CHECK-SAME: (%arg0: tensor<i1>, %arg1: tensor<!tf.resource> {tf.resource_name = "x"})
// CHECK-SAME: (%arg0: tensor<i1>, %arg1: tensor<!tf.resource<tensor<f32>>> {tf_saved_model.bound_input = @x})
// CHECK-NOT: "tf.VarHandleOp"
func @some_args(%arg0: tensor<i1>) {
%0 = "tf.VarHandleOp"() {container = "", shape = "tfshape$", shared_name = "x"} : () -> tensor<!tf.resource>
%0 = "tf.VarHandleOp"() {container = "", shape = "tfshape$", shared_name = "x"} : () -> tensor<!tf.resource<tensor<f32>>>
return
}
// CHECK-LABEL: func @unique_vars
// CHECK-SAME: (%arg0: tensor<!tf.resource<tensor<f32>>> {tf.resource_name = "x"}, %arg1: tensor<!tf.resource<tensor<i32>>> {tf.resource_name = "y"})
// CHECK-SAME: (%arg0: tensor<!tf.resource<tensor<f32>>> {tf_saved_model.bound_input = @x}, %arg1: tensor<!tf.resource<tensor<i32>>> {tf_saved_model.bound_input = @y})
// CHECK-NOT: "tf.VarHandleOp"
func @unique_vars() {
%0 = "tf.VarHandleOp"() {container = "", shape = "tfshape$", shared_name = "x"} : () -> tensor<!tf.resource<tensor<f32>>>
@ -37,7 +41,7 @@ func @unique_vars() {
}
// CHECK-LABEL: func @duplicate_vars
// CHECK-SAME: (%arg0: tensor<!tf.resource<tensor<f32>>> {tf.resource_name = "x"})
// CHECK-SAME: (%arg0: tensor<!tf.resource<tensor<f32>>> {tf_saved_model.bound_input = @x})
// CHECK-NOT: "tf.VarHandleOp"
func @duplicate_vars() {
%0 = "tf.VarHandleOp"() {container = "", shape = "tfshape$", shared_name = "x"} : () -> tensor<!tf.resource<tensor<f32>>>
@ -46,7 +50,7 @@ func @duplicate_vars() {
}
// CHECK-LABEL: func @duplicate_vars_with_users
// CHECK-SAME: (%arg0: tensor<f32>, %arg1: tensor<!tf.resource<tensor<f32>>> {tf.resource_name = "x"})
// CHECK-SAME: (%arg0: tensor<f32>, %arg1: tensor<!tf.resource<tensor<f32>>> {tf_saved_model.bound_input = @x})
// CHECK: "tf.ReadVariableOp"(%arg1)
// CHECK: "tf.AssignAddVariableOp"(%arg1, %arg0)
// CHECK-NOT: "tf.VarHandleOp"

8
tensorflow/compiler/mlir/tensorflow/transforms/passes.h
View File

@ -91,9 +91,11 @@ std::unique_ptr<OperationPass<ModuleOp>> CreateResourceDeviceInferencePass();
// of their aliasing output arguments.
std::unique_ptr<OperationPass<ModuleOp>> CreatePromoteResourcesToArgsPass();
// Creates a pass that promotes tf.VarHandleOp to resource arguments for all
// functions.
std::unique_ptr<OperationPass<ModuleOp>> CreatePromoteVarHandlesToArgsPass();
// Creates a pass that promotes tf.VarHandleOp to to resource arguments of where
// resource names are `tf_saved_model.bound_input` symbol argument attributes
// for all functions.
std::unique_ptr<OperationPass<ModuleOp>>
CreatePromoteVarHandlesToSavedModelArgsPass();
// Creates a pass that converts readonly reference variables to the
// corresponding resource variables.

33
tensorflow/compiler/mlir/tensorflow/transforms/promote_resources_to_args.cc
View File

@ -389,15 +389,18 @@ void PromoteResourcesToArgsPass::runOnOperation() {
return signalPassFailure();
}
class PromoteVarHandlesToArgsPass
: public PassWrapper<PromoteVarHandlesToArgsPass, OperationPass<ModuleOp>> {
// This pass is for promoting Varhandle ops to tf_saved_model.bound_input
// attributes, which are required for TensorFlowSavedModelDialect.
class PromoteVarHandlesToSavedModelArgsPass
: public PassWrapper<PromoteVarHandlesToSavedModelArgsPass,
OperationPass<ModuleOp>> {
public:
void runOnOperation() override;
};
void PromoteVarHandlesToArgsPass::runOnOperation() {
void PromoteVarHandlesToSavedModelArgsPass::runOnOperation() {
ModuleOp module = getOperation();
MLIRContext* context = module.getContext();
for (auto function : module.getOps<FuncOp>()) {
if (failed(CheckSingleBlockFunction(function))) return signalPassFailure();
@ -406,13 +409,15 @@ void PromoteVarHandlesToArgsPass::runOnOperation() {
&var_handle_shared_names);
// Add resource names for each `tf.VarHandleOp` that were promoted to
// resource arguments.
// saved model arguments.
const int var_handle_args_offset =
function.getNumArguments() - var_handle_shared_names.size();
for (auto var_name_and_index : llvm::enumerate(var_handle_shared_names))
for (auto var_name_and_index : llvm::enumerate(var_handle_shared_names)) {
auto symbol_ref =
SymbolRefAttr::get(var_name_and_index.value(), &getContext());
function.setArgAttr(var_name_and_index.index() + var_handle_args_offset,
kResourceNameArgAttr,
StringAttr::get(var_name_and_index.value(), context));
"tf_saved_model.bound_input", symbol_ref);
}
}
}
@ -422,17 +427,19 @@ std::unique_ptr<OperationPass<ModuleOp>> CreatePromoteResourcesToArgsPass() {
return std::make_unique<PromoteResourcesToArgsPass>();
}
std::unique_ptr<OperationPass<ModuleOp>> CreatePromoteVarHandlesToArgsPass() {
return std::make_unique<PromoteVarHandlesToArgsPass>();
std::unique_ptr<OperationPass<ModuleOp>>
CreatePromoteVarHandlesToSavedModelArgsPass() {
return std::make_unique<PromoteVarHandlesToSavedModelArgsPass>();
}
static PassRegistration<PromoteResourcesToArgsPass> pass(
"tf-promote-resources-to-args",
"Promote resources reads/writes to function inputs/outputs.");
static PassRegistration<PromoteVarHandlesToArgsPass> var_handle_pass(
"tf-promote-var-handles-to-args",
"Promote tf.VarHandleOps to function arguments.");
static PassRegistration<PromoteVarHandlesToSavedModelArgsPass> saved_model_pass(
"tf-saved-model-promote-var-handles-to-args",
"Promote tf.VarHandleOps to function arguments in a format of "
"TensorFlowSavedModelDialect.");
} // namespace TF
} // namespace mlir