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: 315275908
Change-Id: Icbc5c032bd9474d279fecf48267665025a53c1bf
This commit is contained in:
Jaesung Chung 2020-06-08 08:04:26 -07:00 committed by TensorFlower Gardener
parent 9429a94225
commit 74fb47ccd2
3 changed files with 24 additions and 37 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,5 +1,4 @@
// 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
// RUN: tf-opt %s -split-input-file -verify-diagnostics -tf-promote-var-handles-to-args | FileCheck %s -dump-input-on-failure
// Tests main function with multiple blocks.
@ -12,27 +11,24 @@ 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<tensor<f32>>> {tf_saved_model.bound_input = @x})
// CHECK-SAME: (%arg0: tensor<!tf.resource> {tf.resource_name = "x"})
// CHECK-NOT: "tf.VarHandleOp"
func @no_args() {
%0 = "tf.VarHandleOp"() {container = "", shape = "tfshape$", shared_name = "x"} : () -> tensor<!tf.resource<tensor<f32>>>
%0 = "tf.VarHandleOp"() {container = "", shape = "tfshape$", shared_name = "x"} : () -> tensor<!tf.resource>
return
}
// CHECK-LABEL: func @some_args
// CHECK-SAME: (%arg0: tensor<i1>, %arg1: tensor<!tf.resource<tensor<f32>>> {tf_saved_model.bound_input = @x})
// CHECK-SAME: (%arg0: tensor<i1>, %arg1: tensor<!tf.resource> {tf.resource_name = "x"})
// CHECK-NOT: "tf.VarHandleOp"
func @some_args(%arg0: tensor<i1>) {
%0 = "tf.VarHandleOp"() {container = "", shape = "tfshape$", shared_name = "x"} : () -> tensor<!tf.resource<tensor<f32>>>
%0 = "tf.VarHandleOp"() {container = "", shape = "tfshape$", shared_name = "x"} : () -> tensor<!tf.resource>
return
}
// CHECK-LABEL: func @unique_vars
// 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-SAME: (%arg0: tensor<!tf.resource<tensor<f32>>> {tf.resource_name = "x"}, %arg1: tensor<!tf.resource<tensor<i32>>> {tf.resource_name = "y"})
// CHECK-NOT: "tf.VarHandleOp"
func @unique_vars() {
%0 = "tf.VarHandleOp"() {container = "", shape = "tfshape$", shared_name = "x"} : () -> tensor<!tf.resource<tensor<f32>>>
@ -41,7 +37,7 @@ func @unique_vars() {
}
// CHECK-LABEL: func @duplicate_vars
// CHECK-SAME: (%arg0: tensor<!tf.resource<tensor<f32>>> {tf_saved_model.bound_input = @x})
// CHECK-SAME: (%arg0: tensor<!tf.resource<tensor<f32>>> {tf.resource_name = "x"})
// CHECK-NOT: "tf.VarHandleOp"
func @duplicate_vars() {
%0 = "tf.VarHandleOp"() {container = "", shape = "tfshape$", shared_name = "x"} : () -> tensor<!tf.resource<tensor<f32>>>
@ -50,7 +46,7 @@ func @duplicate_vars() {
}
// CHECK-LABEL: func @duplicate_vars_with_users
// CHECK-SAME: (%arg0: tensor<f32>, %arg1: tensor<!tf.resource<tensor<f32>>> {tf_saved_model.bound_input = @x})
// CHECK-SAME: (%arg0: tensor<f32>, %arg1: tensor<!tf.resource<tensor<f32>>> {tf.resource_name = "x"})
// CHECK: "tf.ReadVariableOp"(%arg1)
// CHECK: "tf.AssignAddVariableOp"(%arg1, %arg0)
// CHECK-NOT: "tf.VarHandleOp"

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

@ -95,11 +95,9 @@ 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 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 promotes tf.VarHandleOp to resource arguments for all
// functions.
std::unique_ptr<OperationPass<ModuleOp>> CreatePromoteVarHandlesToArgsPass();
// 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,18 +389,15 @@ void PromoteResourcesToArgsPass::runOnOperation() {
return signalPassFailure();
}
// 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>> {
class PromoteVarHandlesToArgsPass
: public PassWrapper<PromoteVarHandlesToArgsPass, OperationPass<ModuleOp>> {
public:
void runOnOperation() override;
};
void PromoteVarHandlesToSavedModelArgsPass::runOnOperation() {
void PromoteVarHandlesToArgsPass::runOnOperation() {
ModuleOp module = getOperation();
MLIRContext* context = module.getContext();
for (auto function : module.getOps<FuncOp>()) {
if (failed(CheckSingleBlockFunction(function))) return signalPassFailure();
@ -409,15 +406,13 @@ void PromoteVarHandlesToSavedModelArgsPass::runOnOperation() {
&var_handle_shared_names);
// Add resource names for each `tf.VarHandleOp` that were promoted to
// saved model arguments.
// resource 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)) {
auto symbol_ref =
SymbolRefAttr::get(var_name_and_index.value(), &getContext());
for (auto var_name_and_index : llvm::enumerate(var_handle_shared_names))
function.setArgAttr(var_name_and_index.index() + var_handle_args_offset,
"tf_saved_model.bound_input", symbol_ref);
}
kResourceNameArgAttr,
StringAttr::get(var_name_and_index.value(), context));
}
}
@ -427,19 +422,17 @@ std::unique_ptr<OperationPass<ModuleOp>> CreatePromoteResourcesToArgsPass() {
return std::make_unique<PromoteResourcesToArgsPass>();
}
std::unique_ptr<OperationPass<ModuleOp>>
CreatePromoteVarHandlesToSavedModelArgsPass() {
return std::make_unique<PromoteVarHandlesToSavedModelArgsPass>();
std::unique_ptr<OperationPass<ModuleOp>> CreatePromoteVarHandlesToArgsPass() {
return std::make_unique<PromoteVarHandlesToArgsPass>();
}
static PassRegistration<PromoteResourcesToArgsPass> pass(
"tf-promote-resources-to-args",
"Promote resources reads/writes to function inputs/outputs.");
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.");
static PassRegistration<PromoteVarHandlesToArgsPass> var_handle_pass(
"tf-promote-var-handles-to-args",
"Promote tf.VarHandleOps to function arguments.");
} // namespace TF
} // namespace mlir