62d50aa8e2
A full legalization conversion stops after the first failed conversion encountered. For building the TF to XLA bridge, it is useful for this pass to continue through and emit information about all of the missing ops. Instead, use the Partial conversion mode to get the full set of operations that are not legalizable. The "full" conversion succeeds if this set is empty. This does not change the behavior when the full legalization pass succeeds. However, if the conversion fails, the outputted error message is now much more useful. For the sake of demonstrating what this might look like with a large model, I've run this on Transformer with the Unary op lowerings removed. Resulting error message output: Before this change: ``` Compilation failure: MLIR TF to XLA legalization failed-:64:11: error: failed to legalize operation 'tf.Rsqrt' -:64:11: note: see current operation: %37 = "tf.Rsqrt"(%33) : (tensor<f32>) -> tensor<f32> ``` After this change (default case): ``` Compilation failure: MLIR TF to XLA legalization failed-:4:3: error: The following operations cannot be legalized: tf.Rsqrt (count: 217); tf.SoftmaxCrossEntropyWithLogits (count: 1); tf.Sqrt (count: 370). These legalization failure(s) may be due to missing TF to HLO lowerings and/or unsupported attributes, etc. -:4:3: error: Emitting more detail about one op that failed to legalize... -:251:12: error: 'tf.Rsqrt' op is not legalizable -:251:12: note: see current operation: %224 = "tf.Rsqrt"(%220) : (tensor<f32>) -> tensor<f32> ``` After this change (verbose case, with logging set to 1): ``` Compilation failure: MLIR TF to XLA legalization failed-:4:3: error: The following operations cannot be legalized: tf.Rsqrt (count: 217); tf.SoftmaxCrossEntropyWithLogits (count: 1); tf.Sqrt (count: 370). These legalization failure(s) may be due to missing TF to HLO lowerings and/or unsupported attributes, etc. -:4:3: error: Emitting more detail about one of each type of op that failed to legalize... -:1769:13: error: 'tf.Rsqrt' op is not legalizable -:1769:13: note: see current operation: %1742 = "tf.Rsqrt"(%1738) : (tensor<f32>) -> tensor<f32> -:3308:24: error: 'tf.SoftmaxCrossEntropyWithLogits' op is not legalizable -:3308:24: note: see current operation: %loss, %backprop = "tf.SoftmaxCrossEntropyWithLogits"(%3495, %3503) : (tensor<768x33708xf32>, tensor<768x33708xf32>) -> (tensor<768xf32>, tensor<768x33708xf32>) -:6944:13: error: 'tf.Sqrt' op is not legalizable -:6944:13: note: see current operation: %7319 = "tf.Sqrt"(%7318) : (tensor<f32>) -> tensor<f32> ``` PiperOrigin-RevId: 310388184 Change-Id: Idea12490be1ace18cfb89bda51f5f948bd2d691c
45 lines
2.1 KiB
MLIR
45 lines
2.1 KiB
MLIR
// RUN: tf-opt %s -xla-legalize-tf -split-input-file -verify-diagnostics
|
|
|
|
// expected-error@below{{The following operations cannot be legalized: tf.NoOp (count: 1); tf_executor.fetch (count: 1); tf_executor.graph (count: 1); tf_executor.island (count: 1); tf_executor.yield (count: 1). These legalization failure(s) may be due to missing TF to HLO lowerings and/or unsupported attributes, etc.}}
|
|
// expected-error@below{{Emitting more detail about one op that failed to legalize...}}
|
|
func @tf_executor_graph_op() {
|
|
tf_executor.graph {
|
|
%0 = tf_executor.island {
|
|
// expected-error@+1 {{'tf.NoOp' op is not legalizable}}
|
|
"tf.NoOp"() {} : () -> ()
|
|
tf_executor.yield
|
|
}
|
|
tf_executor.fetch
|
|
}
|
|
return
|
|
}
|
|
|
|
// -----
|
|
|
|
// expected-error@below{{The following operations cannot be legalized: tf.OpA (count: 1). These legalization failure(s) may be due to missing TF to HLO lowerings and/or unsupported attributes, etc.}}
|
|
func @tf_unknown_op(%arg0: tensor<2xi32>) -> tensor<2xi32> {
|
|
// expected-error@+1 {{'tf.OpA' op is not legalizable}}
|
|
%0 = "tf.OpA"(%arg0, %arg0) : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi32>
|
|
return %0: tensor<2xi32>
|
|
}
|
|
|
|
// -----
|
|
|
|
func @tf_known_op(%arg0: tensor<2xi32>) -> tensor<2xi32> {
|
|
%0 = "tf.Add"(%arg0, %arg0) : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi32>
|
|
return %0: tensor<2xi32>
|
|
}
|
|
|
|
// -----
|
|
|
|
// expected-error@below{{The following operations cannot be legalized: tf.OpA (count: 1); tf.OpB (count: 2). These legalization failure(s) may be due to missing TF to HLO lowerings and/or unsupported attributes, etc.}}
|
|
// expected-error@below{{Emitting more detail about one op that failed to legalize...}}
|
|
func @tf_unknown_known_mix(%arg0: tensor<2xi32>) -> tensor<2xi32> {
|
|
// expected-error@+1 {{'tf.OpA' op is not legalizable}}
|
|
%0 = "tf.OpA"(%arg0, %arg0) : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi32>
|
|
%1 = "tf.OpB"(%0, %0) : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi32>
|
|
%2 = "tf.Add"(%1, %1) : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi32>
|
|
%3 = "tf.OpB"(%2, %2) : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi32>
|
|
return %2: tensor<2xi32>
|
|
}
|