From cfa9e61e4d296b3ecd4a5fc0bac8cf372698bd44 Mon Sep 17 00:00:00 2001
From: Christian Sigg <csigg@google.com>
Date: Thu, 14 Jan 2021 12:11:04 -0800
Subject: [PATCH] Use mlir::OpState::operator->() to get to methods of
 mlir::Operation.

This is a preparation step to remove those methods from OpState.

PiperOrigin-RevId: 351850594
Change-Id: I5ba53e37f283b8f5e028be043119859c487097e4
---
 .../lite/utils/perception_ops_utils_test.cc   |  2 +-
 .../transforms/tpu_resource_partitioning.cc   |  4 +--
 .../utils/tpu_rewrite_device_util_test.cc     | 34 +++++++++----------
 3 files changed, 20 insertions(+), 20 deletions(-)

diff --git a/tensorflow/compiler/mlir/lite/utils/perception_ops_utils_test.cc b/tensorflow/compiler/mlir/lite/utils/perception_ops_utils_test.cc
index 19a2b81325c..cb3a4fc9af5 100644
--- a/tensorflow/compiler/mlir/lite/utils/perception_ops_utils_test.cc
+++ b/tensorflow/compiler/mlir/lite/utils/perception_ops_utils_test.cc
@@ -44,7 +44,7 @@ FuncOp createMaxUnpoolingFunc(
 
   func.addEntryBlock();
   mlir::StringAttr attr_value = builder->getStringAttr("MaxUnpooling2D");
-  func.setAttr("tf._implements", attr_value);
+  func->setAttr("tf._implements", attr_value);
   return func;
 }
 
diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/tpu_resource_partitioning.cc b/tensorflow/compiler/mlir/tensorflow/transforms/tpu_resource_partitioning.cc
index e094603f8eb..4c7936b74b6 100644
--- a/tensorflow/compiler/mlir/tensorflow/transforms/tpu_resource_partitioning.cc
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/tpu_resource_partitioning.cc
@@ -73,8 +73,8 @@ Type GetResourceSubtype(Value resource) {
 // resource reads and writes per associated partitioned resource handle.
 void PartitionResourceReadsWrites(tf_device::ClusterFuncOp cluster_func) {
   bool use_spmd = false;
-  if (auto use_spmd_attr =
-          cluster_func.getAttrOfType<BoolAttr>("use_spmd_for_xla_partitioning"))
+  if (auto use_spmd_attr = cluster_func->getAttrOfType<BoolAttr>(
+          "use_spmd_for_xla_partitioning"))
     use_spmd = use_spmd_attr.getValue();
 
   if (!use_spmd) return;
diff --git a/tensorflow/compiler/mlir/tensorflow/utils/tpu_rewrite_device_util_test.cc b/tensorflow/compiler/mlir/tensorflow/utils/tpu_rewrite_device_util_test.cc
index 63afd0a93e6..0f211512b0a 100644
--- a/tensorflow/compiler/mlir/tensorflow/utils/tpu_rewrite_device_util_test.cc
+++ b/tensorflow/compiler/mlir/tensorflow/utils/tpu_rewrite_device_util_test.cc
@@ -691,9 +691,9 @@ TEST(TPURewriteDeviceUtilTest, TestGetHostDeviceFailMissingDeviceAssignment) {
   llvm::SmallVector<mlir::Type, 8> result_types;
   auto cluster = builder.create<mlir::tf_device::ClusterOp>(
       mlir::UnknownLoc::get(&context), result_types);
-  cluster.setAttr(kNumCoresPerReplicaAttr,
-                  builder.getIntegerAttr(builder.getIntegerType(64), 1));
-  cluster.setAttr(kTopologyAttr, builder.getStringAttr(""));
+  cluster->setAttr(kNumCoresPerReplicaAttr,
+                   builder.getIntegerAttr(builder.getIntegerType(64), 1));
+  cluster->setAttr(kTopologyAttr, builder.getStringAttr(""));
 
   mlir::TF::RuntimeDevices runtime_devices;
   std::string host_device;
@@ -711,12 +711,12 @@ TEST(TPURewriteDeviceUtilTest, TestGetHostDeviceFailBadDeviceAssignment) {
   llvm::SmallVector<mlir::Type, 8> result_types;
   auto cluster = builder.create<mlir::tf_device::ClusterOp>(
       mlir::UnknownLoc::get(&context), result_types);
-  cluster.setAttr(kNumCoresPerReplicaAttr,
-                  builder.getIntegerAttr(builder.getIntegerType(64), 1));
-  cluster.setAttr(kTopologyAttr, builder.getStringAttr(""));
-  cluster.setAttr(kDeviceAssignmentAttr,
-                  builder.getStrArrayAttr(llvm::ArrayRef<llvm::StringRef>(
-                      {"bad_device_assigment"})));
+  cluster->setAttr(kNumCoresPerReplicaAttr,
+                   builder.getIntegerAttr(builder.getIntegerType(64), 1));
+  cluster->setAttr(kTopologyAttr, builder.getStringAttr(""));
+  cluster->setAttr(kDeviceAssignmentAttr,
+                   builder.getStrArrayAttr(llvm::ArrayRef<llvm::StringRef>(
+                       {"bad_device_assigment"})));
 
   mlir::TF::RuntimeDevices runtime_devices;
   std::string host_device;
@@ -737,10 +737,10 @@ TEST(TPURewriteDeviceUtilTest, TestGetHostDeviceFailBadDeviceName) {
   llvm::SmallVector<mlir::Type, 8> result_types;
   auto cluster = builder.create<mlir::tf_device::ClusterOp>(
       mlir::UnknownLoc::get(&context), result_types);
-  cluster.setAttr(kNumCoresPerReplicaAttr,
-                  builder.getIntegerAttr(builder.getIntegerType(64), 1));
-  cluster.setAttr(kTopologyAttr, builder.getStringAttr(""));
-  cluster.setAttr(kDeviceAssignmentAttr, builder.getArrayAttr({}));
+  cluster->setAttr(kNumCoresPerReplicaAttr,
+                   builder.getIntegerAttr(builder.getIntegerType(64), 1));
+  cluster->setAttr(kTopologyAttr, builder.getStringAttr(""));
+  cluster->setAttr(kDeviceAssignmentAttr, builder.getArrayAttr({}));
 
   mlir::TF::RuntimeDevices runtime_devices;
   GetDevicesFromOp(*module_ref, &runtime_devices);
@@ -791,10 +791,10 @@ TEST(TPURewriteDeviceUtilTest, TestGetHostDeviceNotReplicated) {
   llvm::SmallVector<mlir::Type, 8> result_types;
   auto cluster = builder.create<mlir::tf_device::ClusterOp>(
       mlir::UnknownLoc::get(&context), result_types);
-  cluster.setAttr(kNumCoresPerReplicaAttr,
-                  builder.getIntegerAttr(builder.getIntegerType(64), 1));
-  cluster.setAttr(kTopologyAttr, builder.getStringAttr(""));
-  cluster.setAttr(kDeviceAssignmentAttr, builder.getArrayAttr({}));
+  cluster->setAttr(kNumCoresPerReplicaAttr,
+                   builder.getIntegerAttr(builder.getIntegerType(64), 1));
+  cluster->setAttr(kTopologyAttr, builder.getStringAttr(""));
+  cluster->setAttr(kDeviceAssignmentAttr, builder.getArrayAttr({}));
 
   mlir::TF::RuntimeDevices runtime_devices;
   GetDevicesFromOp(*module_ref, &runtime_devices);