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
This commit is contained in:
Christian Sigg 2021-01-14 12:11:04 -08:00 committed by TensorFlower Gardener
parent 07592d94b6
commit cfa9e61e4d
3 changed files with 20 additions and 20 deletions

View File

@ -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;
}

View File

@ -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;

View File

@ -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);