[MLIR]The following operations cannot be legalized: tf.VariableV2

TensorFlow Special Interest Groups MLIR
1

Hello, this is my first time posting here. Let me know whether this question is suited to be posted to this forum or not.

I have been trying to translate tensorflow models into MLIR-HLO. As such I successfully used tf-mlir-translate , tf-opt with very simple TF models like Conv2D, MatMul,MaxPool…, achieving the expected results.

However, if I try to do the same with the official example tf.Variable with the three steps below:

tf-mlir-translate -graphdef-to-mlir -tf-enable-shape-inference-on-import=false target.pbtxt -tf-prune-unused-nodes -tf-control-output-arrays=Variable/Assign,AssignAdd -o target.mlir
tf-opt -tf-executor-to-functional-conversion target.mlir -o target-func.mlir
tf-opt --tf-to-hlo-pipeline target-func.mlir -o target-mhlo.mlir

I get the following error message with the 3’rd step:

target-func.mlir:2:3: error: The following operations cannot be legalized: tf.Assign (count: 1); tf.AssignAdd (count: 1); tf.VariableV2 (count: 1). These legalization failure(s) may be due to missing TF to HLO lowerings and/or unsupported attributes, etc.
  func @main() attributes {tf.entry_function = {control_outputs = "Variable/Assign,AssignAdd", inputs = "", outputs = ""}} {
  ^
target-func.mlir:2:3: error: Emitting more detail about one op that failed to legalize...
  func @main() attributes {tf.entry_function = {control_outputs = "Variable/Assign,AssignAdd", inputs = "", outputs = ""}} {
  ^
target-func.mlir:7:10: error: 'tf.Assign' op is not legalizable
    %2 = "tf.Assign"(%0, %cst_0) {_class = ["loc:@Variable"], device = "", use_locking = true, validate_shape = true} : (tensor<!tf.f32ref>, tensor<f32>) -> tensor<*x!tf.f32ref>
         ^
target-func.mlir:7:10: note: see current operation: %4 = "tf.Assign"(%2, %0) {_class = ["loc:@Variable"], device = "", use_locking = true, validate_shape = true} : (tensor<!tf.f32ref>, tensor<f32>) -> tensor<!tf.f32ref>

As of right now, is there a way to lower these tf operations into mhlo-hlo ?

1 Like
3

“HLO ops are intended for the numeric computation side and are mostly pure with state management is handled outside them. For TF this is done via TF ops and the state hoisted/sunk outside stateless parts"

How to continue lowering for TF? does TF lower the resource variables to other dialect than mhlo-hlo?

1 Like
5

In general (for models that are supporting it), the resource variables will be hoisted outside of the function intended to be compiler to make it “pure”.

See for example this pass pipeline which is our flow to target TPUs: tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/bridge.cc at master · tensorflow/tensorflow · GitHub

7

Hi GipsonLeo,

Have you found any solutions to convert the tf.assign to MHLO dialect? I experienced the same issue.

Thanks

1 Like