Promote resources to args with tf-opt

Hello,

I am trying to apply the command tf-opt --promote-resources-to-args on the following TF/MLIR code.

func @counter(%arg0: tensor) → tensor {
%1 = “tf.VarHandleOp”() {container = “”, shared_name = “x”} : () → tensor<!tf_type.resource<tensor>>
%2 = “tf.ReadVariableOp”(%1) : (tensor<!tf_type.resource<tensor>>) → tensor
%3 = “tf.Add”(%arg0, %2): (tensor, tensor) → tensor
“tf.AssignVariableOp”(%1, %3) {device = “”} : (tensor<!tf_type.resource<tensor>>, tensor) → ()
%4 = “tf.ReadVariableOp”(%1) : (tensor<!tf_type.resource<tensor>>) → tensor
return %4: tensor
}

It does nothing because the function’s name is not @main. Is there a way to apply this command on every function whatever its name is ?

The result I need is :

func @counter(%arg0: tensor, %arg1: tensor {tf.aliasing_output = 1 : i64, tf.resource_name = “x”}) → (tensor, tensor) {
%0 = “tf.Add”(%arg0, %arg1) : (tensor, tensor) → tensor
return %0, %0 : tensor, tensor
}

Best regards,

HP

Unfortunately this isn’t possible as-is since it is hard-coded here: https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/mlir/tensorflow/transforms/promote_resources_to_args.cc#L348-L349

Feel free to send a patch that would make it an option of the pass to specify the entrypoint name!

1 Like

Thank you for your answer !