Change promote-resources-to-args pass to emit all the unexpected users for
VarHandleOp.
Modify an existing test to test the output.
PiperOrigin-RevId: 302350400
Change-Id: I0cb30a477c13629b20172b77e1e6340392f8fd5b
diff --git a/tensorflow/compiler/mlir/tensorflow/tests/promote_resources_to_args.mlir b/tensorflow/compiler/mlir/tensorflow/tests/promote_resources_to_args.mlir
index dcf2127..db15aeb 100644
--- a/tensorflow/compiler/mlir/tensorflow/tests/promote_resources_to_args.mlir
+++ b/tensorflow/compiler/mlir/tensorflow/tests/promote_resources_to_args.mlir
@@ -297,8 +297,9 @@
// Tests VarHandleOp has users that are not removed.
func @main() -> tensor<i1> {
- // expected-error@+1 {{expects no uses}}
+ // expected-error@+1 {{expects no uses but used by operations: tf.UnknownOp, tf.VarIsInitializedOp}}
%0 = "tf.VarHandleOp"() {container = "", shape = "tfshape$", shared_name = "x"} : () -> tensor<!tf.resource<tensor<f32>>>
%1 = "tf.VarIsInitializedOp"(%0) : (tensor<!tf.resource<tensor<f32>>>) -> tensor<i1>
+ %2 = "tf.UnknownOp"(%0) : (tensor<!tf.resource<tensor<f32>>>) -> tensor<i1>
return %1 : tensor<i1>
}
diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/promote_resources_to_args.cc b/tensorflow/compiler/mlir/tensorflow/transforms/promote_resources_to_args.cc
index e69ac6a..a0a5217 100644
--- a/tensorflow/compiler/mlir/tensorflow/transforms/promote_resources_to_args.cc
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/promote_resources_to_args.cc
@@ -36,6 +36,7 @@
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/PointerUnion.h"
#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallSet.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/Casting.h"
#include "mlir/Dialect/StandardOps/IR/Ops.h" // TF:llvm-project
@@ -210,8 +211,21 @@
for (Operation& op : llvm::make_early_inc_range(function.front())) {
auto var_handle_op = llvm::dyn_cast<TF::VarHandleOp>(op);
if (!var_handle_op) continue;
- if (!var_handle_op.use_empty())
- return var_handle_op.emitOpError() << "expects no uses";
+ if (!var_handle_op.use_empty()) {
+ // SmallSet will use a vector when there is only one element and use
+ // std::set when there are more than one elements. This ensures that
+ // the operations in the error message are ordered.
+ llvm::SmallSet<std::string, 2> unique_operations;
+ llvm::for_each(
+ var_handle_op.getOperation()->getUsers(), [&](Operation* user) {
+ unique_operations.insert(user->getName().getStringRef().str());
+ });
+
+ return var_handle_op.emitOpError(
+ "expects no uses but used by operations: ")
+ << llvm::join(unique_operations.begin(), unique_operations.end(),
+ ", ");
+ }
op.erase();
}
