From 251923169ddde0b69c58fb41b23443423ababd2d Mon Sep 17 00:00:00 2001 From: Rahul Joshi Date: Mon, 22 Jun 2020 09:58:59 -0700 Subject: [PATCH] [NFC] Eliminate use of .getBlocks() when not needed Also use llvm::hasSingleElement() instead of .size() == 1 PiperOrigin-RevId: 317675565 Change-Id: I4f0e8892957c2b20e115584fe7424da68d53b67a --- tensorflow/compiler/mlir/lite/flatbuffer_export.cc | 6 +++--- .../compiler/mlir/lite/transforms/post_quantize.cc | 2 +- tensorflow/compiler/mlir/tensorflow/ir/tf_device.cc | 6 +++--- tensorflow/compiler/mlir/tensorflow/ir/tf_executor.cc | 6 +++--- .../mlir/tensorflow/transforms/fold_switch.cc | 2 +- .../transforms/materialize_mlir_passthrough_op.cc | 2 +- .../transforms/promote_resources_to_args.cc | 4 ++-- .../mlir/tensorflow/transforms/resource_op_lifting.cc | 2 +- .../transforms/tpu_sharding_identification_pass.cc | 11 ++++------- .../mlir/tensorflow/translate/export_graphdef.cc | 2 +- .../tensorflow/translate/tf_functional_to_executor.cc | 4 ++-- .../tensorflow/translate/translate_tf_dialect_op.cc | 5 +++-- tensorflow/compiler/mlir/xla/mlir_hlo_to_hlo.cc | 4 ++-- .../mlir/xla/transforms/hlo_legalize_to_lhlo.cc | 8 ++++---- .../mlir/xla/transforms/legalize_tf_with_tf2xla.cc | 8 +++++--- .../compiler/mlir/xla/transforms/lhlo_fuse_linalg.cc | 5 +++-- .../compiler/mlir/xla/transforms/xla_hlo_fusion.cc | 2 +- .../compiler/xla/service/mlir_gpu/kernel_lowering.cc | 2 +- 18 files changed, 41 insertions(+), 40 deletions(-) diff --git a/tensorflow/compiler/mlir/lite/flatbuffer_export.cc b/tensorflow/compiler/mlir/lite/flatbuffer_export.cc index a260670015a..e34e7ae7ca6 100644 --- a/tensorflow/compiler/mlir/lite/flatbuffer_export.cc +++ b/tensorflow/compiler/mlir/lite/flatbuffer_export.cc @@ -240,10 +240,10 @@ static bool IsValidTFLiteMlirModule(ModuleOp module) { } for (auto fn : module.getOps()) { - if (fn.getBlocks().size() != 1) { + if (!llvm::hasSingleElement(fn)) { return fn.emitError("should have exactly one basic block"), false; } - auto& bb = fn.getBlocks().front(); + auto& bb = fn.front(); for (auto arg : bb.getArguments()) { if (!HasValidTFLiteType(arg, fn)) @@ -1089,7 +1089,7 @@ void Translator::InitializeNamesFromAttribute(FuncOp fn, bool* has_input_attr) { dict_attr.get("outputs").dyn_cast_or_null()) { str.getValue().split(output_names, ',', /*MaxSplit=*/-1, /*KeepEmpty=*/false); - auto term = fn.getBlocks().back().getTerminator(); + auto term = fn.back().getTerminator(); if (output_names.size() != term->getNumOperands()) { fn.emitWarning() << "output names (" << output_names.size() << ") != terminator operands (" << term->getNumOperands() diff --git a/tensorflow/compiler/mlir/lite/transforms/post_quantize.cc b/tensorflow/compiler/mlir/lite/transforms/post_quantize.cc index 9a1da0ad03d..33380e00543 100644 --- a/tensorflow/compiler/mlir/lite/transforms/post_quantize.cc +++ b/tensorflow/compiler/mlir/lite/transforms/post_quantize.cc @@ -52,7 +52,7 @@ class PostQuantizePass : public PassWrapper { void RemoveQuantizationAdaptorOps(FuncOp func) { mlir::OpBuilder builder(func.getBody()); - auto& bb = func.getBlocks().front(); + auto& bb = func.front(); auto* terminator = bb.getTerminator(); int num_args = bb.getNumArguments(); diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_device.cc b/tensorflow/compiler/mlir/tensorflow/ir/tf_device.cc index b8f0585040c..7dd74282487 100644 --- a/tensorflow/compiler/mlir/tensorflow/ir/tf_device.cc +++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_device.cc @@ -299,13 +299,13 @@ ParseResult ParseReplicateOp(OpAsmParser* parser, OperationState* state) { parser->parseRegion(body, region_args, region_arg_types)) return failure(); - if (body.getBlocks().size() > 1) - return parser->emitError(loc) << "expects a single block region"; - // Ensure that the region is well formed: it contains at least a block with // a ReturnOp terminator. ReplicateOp::ensureTerminator(body, parser->getBuilder(), state->location); + if (!llvm::hasSingleElement(body)) + return parser->emitError(loc) << "expects a single block region"; + Operation& terminator = body.front().back(); if (!isa(terminator)) return parser->emitError(loc) << "expects a tf_device.return terminator"; diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_executor.cc b/tensorflow/compiler/mlir/tensorflow/ir/tf_executor.cc index 3403651eef8..1e66eee06bb 100644 --- a/tensorflow/compiler/mlir/tensorflow/ir/tf_executor.cc +++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_executor.cc @@ -220,13 +220,13 @@ ParseResult ParseGraphOp(OpAsmParser &parser, OperationState &result) { Region &body = *result.addRegion(); if (parser.parseRegion(body, llvm::None, llvm::None)) return failure(); - if (body.getBlocks().size() > 1) - return parser.emitError(loc) << "expects a single block region"; - // Ensure that the region is well formed: it contains at least a block with // a FetchOp terminator. GraphOp::ensureTerminator(body, parser.getBuilder(), result.location); + if (!llvm::hasSingleElement(body)) + return parser.emitError(loc) << "expects a single block region"; + // Get the results type from the terminator type inside the graph. Operation &fetch = body.back().back(); if (!isa(fetch)) diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/fold_switch.cc b/tensorflow/compiler/mlir/tensorflow/transforms/fold_switch.cc index 4d26747ebdc..b47378762a9 100644 --- a/tensorflow/compiler/mlir/tensorflow/transforms/fold_switch.cc +++ b/tensorflow/compiler/mlir/tensorflow/transforms/fold_switch.cc @@ -199,7 +199,7 @@ static void MatchSwitchFoldOps(tf_executor::SwitchOp switch_op, // Folds merge nodes with only a single non-dead input. static LogicalResult FoldMergeNodes(FuncOp function, const DeadQueue& queue) { // Create builder for val_index of MergeOp. - auto* block = &function.getBlocks().front(); + auto* block = &function.front(); OpBuilder builder = OpBuilder::atBlockEnd(block); auto type = builder.getIntegerType(32); auto build_index = [&](Location loc, int value) { diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/materialize_mlir_passthrough_op.cc b/tensorflow/compiler/mlir/tensorflow/transforms/materialize_mlir_passthrough_op.cc index 94fdfb310ac..3ed27d7ce30 100644 --- a/tensorflow/compiler/mlir/tensorflow/transforms/materialize_mlir_passthrough_op.cc +++ b/tensorflow/compiler/mlir/tensorflow/transforms/materialize_mlir_passthrough_op.cc @@ -71,7 +71,7 @@ void MaterializePassthroughOpPass::runOnFunction() { return; } Region &body = main.getBody(); - if (body.getBlocks().size() != 1) { + if (!llvm::hasSingleElement(body)) { op->emitError() << "MLIR Opaque Op expects a main() entry point with a " "single block\n"; return; 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 cece23b4750..af36770f496 100644 --- a/tensorflow/compiler/mlir/tensorflow/transforms/promote_resources_to_args.cc +++ b/tensorflow/compiler/mlir/tensorflow/transforms/promote_resources_to_args.cc @@ -80,11 +80,11 @@ constexpr char kResourceNameArgAttr[] = "tf.resource_name"; // Checks if a function has only one block. mlir::LogicalResult CheckSingleBlockFunction(FuncOp function) { - if (!hasSingleElement(function.getBlocks())) + if (!llvm::hasSingleElement(function)) { return function.emitError() << "expects function '" << function.getName() << "' to have 1 block, got " << function.getBlocks().size(); - + } return success(); } diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/resource_op_lifting.cc b/tensorflow/compiler/mlir/tensorflow/transforms/resource_op_lifting.cc index ed7ebc25c9f..799ab3a0f0d 100644 --- a/tensorflow/compiler/mlir/tensorflow/transforms/resource_op_lifting.cc +++ b/tensorflow/compiler/mlir/tensorflow/transforms/resource_op_lifting.cc @@ -1113,7 +1113,7 @@ LogicalResult ResourceLiftingForFunctionalControlFlow(FuncOp function) { // This routine should only be called when control flow operations are still // represented with TF IfOp and WhileOp operations. In this case, there should // be only one basic blocks in the MLIR representation. - if (!hasSingleElement(function.getBlocks())) { + if (!llvm::hasSingleElement(function)) { return function.emitError() << "expect the function to have 1 block while it has " << function.getBlocks().size(); diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/tpu_sharding_identification_pass.cc b/tensorflow/compiler/mlir/tensorflow/transforms/tpu_sharding_identification_pass.cc index f8b6e364f55..b05e87c6485 100644 --- a/tensorflow/compiler/mlir/tensorflow/transforms/tpu_sharding_identification_pass.cc +++ b/tensorflow/compiler/mlir/tensorflow/transforms/tpu_sharding_identification_pass.cc @@ -159,8 +159,7 @@ llvm::SmallVector ExtractFunctionsConnectedToArg( while (!functions_to_parse.empty()) { llvm::SmallVector newly_discovered_functions; for (auto function_info : functions_to_parse) { - Block& func_entry_block = - function_info.func.getBody().getBlocks().front(); + Block& func_entry_block = function_info.func.front(); auto argument = func_entry_block.getArgument(function_info.argument_index); @@ -186,8 +185,7 @@ void IdentifyXlaShardingForComputationInputs( StringRef logical_core_0_sharding, tf_device::ClusterFuncOp cluster_func_op, FuncOp cluster_function, Builder* builder) { // Look up function definition from module. - Block& cluster_function_block = - cluster_function.getBody().getBlocks().front(); + Block& cluster_function_block = cluster_function.front(); ModuleOp module = cluster_func_op.getParentOfType(); llvm::SmallVector sharding_for_args( @@ -215,8 +213,7 @@ void IdentifyXlaShardingForComputationInputs( const int function_argument_index = function_arg_info.argument_index; auto& parsed_function = function_arg_info.func; - Block& parsed_function_block = - parsed_function.getBody().getBlocks().front(); + Block& parsed_function_block = parsed_function.front(); arg_sharding = ParseInputSharding( parsed_function_block.getArgument(function_argument_index)); } @@ -245,7 +242,7 @@ void IdentifyXlaShardingForComputationOutputs( tf_device::ClusterFuncOp cluster_func, Builder* builder) { // By default return values from logical core 0 is used if no sharding // configuration is defined. - Block& function_block = func.getBody().getBlocks().front(); + Block& function_block = func.front(); Operation* terminator = function_block.getTerminator(); llvm::SmallVector sharding_for_rets( terminator->getNumOperands(), logical_core_0_sharding); diff --git a/tensorflow/compiler/mlir/tensorflow/translate/export_graphdef.cc b/tensorflow/compiler/mlir/tensorflow/translate/export_graphdef.cc index 262f6f4e50c..8cd14894f8f 100644 --- a/tensorflow/compiler/mlir/tensorflow/translate/export_graphdef.cc +++ b/tensorflow/compiler/mlir/tensorflow/translate/export_graphdef.cc @@ -128,7 +128,7 @@ class LegalizedOpOrValLocNameMapper : public OpOrArgLocNameMapper { Status HasSingleGraphSingleOpIslandsFunctions(mlir::ModuleOp module) { Status status = Status::OK(); module.walk([&](mlir::FuncOp function) { - if (function.getBlocks().size() != 1) { + if (!llvm::hasSingleElement(function)) { status = errors::FailedPrecondition( kInvalidExecutorGraphMsg, "only single block functions are supported."); diff --git a/tensorflow/compiler/mlir/tensorflow/translate/tf_functional_to_executor.cc b/tensorflow/compiler/mlir/tensorflow/translate/tf_functional_to_executor.cc index 29f98de6448..78019119d9d 100644 --- a/tensorflow/compiler/mlir/tensorflow/translate/tf_functional_to_executor.cc +++ b/tensorflow/compiler/mlir/tensorflow/translate/tf_functional_to_executor.cc @@ -46,13 +46,13 @@ struct FunctionalToExecutorDialectConversion } // end anonymous namespace void FunctionalToExecutorDialectConversion::runOnFunction() { - if (getFunction().getBlocks().size() != 1) { + if (!llvm::hasSingleElement(getFunction())) { LLVM_DEBUG(llvm::dbgs() << "Expect single block function, skip conversion " "to tf_executor dialect\n"); return; } auto loc = getFunction().getLoc(); - mlir::Block& body = getFunction().getBody().front(); + mlir::Block& body = getFunction().front(); // Find region of interest and ReturnOp. auto copy_range = body.without_terminator(); if (copy_range.begin() != copy_range.end() && diff --git a/tensorflow/compiler/mlir/tensorflow/translate/translate_tf_dialect_op.cc b/tensorflow/compiler/mlir/tensorflow/translate/translate_tf_dialect_op.cc index bd3fe9876ff..5236bdeffbf 100644 --- a/tensorflow/compiler/mlir/tensorflow/translate/translate_tf_dialect_op.cc +++ b/tensorflow/compiler/mlir/tensorflow/translate/translate_tf_dialect_op.cc @@ -13,6 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#include "llvm/ADT/STLExtras.h" #include "llvm/Support/ToolOutputFile.h" #include "mlir/IR/Function.h" // from @llvm-project #include "mlir/IR/Location.h" // from @llvm-project @@ -26,12 +27,12 @@ static mlir::Operation* ExtractOnlyOp(mlir::ModuleOp module) { mlir::FuncOp fn = module.lookupSymbol("main"); if (!fn) return nullptr; - if (fn.getBlocks().size() != 1) return nullptr; + if (!llvm::hasSingleElement(fn)) return nullptr; // Here, modules with exactly two operations in the only basic block are // supported. The last operation should be a terminator operation and the // other operation is the operation of interest. - auto& block = fn.getBlocks().front(); + auto& block = fn.front(); if (block.getOperations().size() != 2) return nullptr; if (!block.back().isKnownTerminator()) return nullptr; diff --git a/tensorflow/compiler/mlir/xla/mlir_hlo_to_hlo.cc b/tensorflow/compiler/mlir/xla/mlir_hlo_to_hlo.cc index 60d9a698731..7a576780c61 100644 --- a/tensorflow/compiler/mlir/xla/mlir_hlo_to_hlo.cc +++ b/tensorflow/compiler/mlir/xla/mlir_hlo_to_hlo.cc @@ -1148,13 +1148,13 @@ LogicalResult ConvertToHloModule::LowerFunctionCall( LogicalResult ConvertToHloModule::RunOnFunction(mlir::FuncOp f) { if (lowered_computation_.count(f)) return success(); - if (f.getBlocks().size() != 1) { + if (!llvm::hasSingleElement(f)) { return f.emitError("only single block Function supported"); } // Create a sub-builder if this is not the main function. std::unique_ptr builder_up; - bool entry_function = f.getName().str() == "main"; + bool entry_function = f.getName() == "main"; if (!entry_function) builder_up = module_builder_.CreateSubBuilder(f.getName().str()); auto& builder = entry_function ? module_builder_ : *builder_up; diff --git a/tensorflow/compiler/mlir/xla/transforms/hlo_legalize_to_lhlo.cc b/tensorflow/compiler/mlir/xla/transforms/hlo_legalize_to_lhlo.cc index 446f2aae833..7cdc0d92207 100644 --- a/tensorflow/compiler/mlir/xla/transforms/hlo_legalize_to_lhlo.cc +++ b/tensorflow/compiler/mlir/xla/transforms/hlo_legalize_to_lhlo.cc @@ -230,10 +230,10 @@ struct HloToLhloReduceOpConverter : public BaseOpConversion { auto loc = op.getLoc(); // TODO(b/137624192) Implement variadic reduce. if (op.getNumResults() != 1) return failure(); - if (op.getParentRegion()->getBlocks().size() != 1) { - op.emitOpError() << "tensor to buffer conversion expects a single block " - "in the region containing the operation"; - return failure(); + if (!llvm::hasSingleElement(op.body())) { + return op.emitOpError() + << "tensor to buffer conversion expects a single block " + "in the region containing the operation"; } const auto& original_results = op.getResults(); SmallVector buffer_args(operands.begin(), operands.end()); diff --git a/tensorflow/compiler/mlir/xla/transforms/legalize_tf_with_tf2xla.cc b/tensorflow/compiler/mlir/xla/transforms/legalize_tf_with_tf2xla.cc index 8f96f4d1305..54453406ef7 100644 --- a/tensorflow/compiler/mlir/xla/transforms/legalize_tf_with_tf2xla.cc +++ b/tensorflow/compiler/mlir/xla/transforms/legalize_tf_with_tf2xla.cc @@ -23,6 +23,7 @@ limitations under the License. #include "absl/strings/string_view.h" #include "llvm/ADT/DenseSet.h" #include "llvm/ADT/Optional.h" +#include "llvm/ADT/STLExtras.h" #include "mlir/Dialect/StandardOps/IR/Ops.h" // from @llvm-project #include "mlir/IR/Diagnostics.h" // from @llvm-project #include "mlir/IR/Function.h" // from @llvm-project @@ -320,13 +321,14 @@ LogicalResult FuncLegalizer::PrepareParams() { } LogicalResult FuncLegalizer::Legalize() { + if (func_.empty()) return success(); + // TensorFlow functions don't use CFGs. - if (func_.getBlocks().size() > 1) { + if (!llvm::hasSingleElement(func_)) { emitError(func_.getLoc()) << "requires at most one block in a TF function"; return failure(); } - if (func_.getBlocks().empty()) return success(); - Block& block = func_.getBlocks().front(); + Block& block = func_.front(); std::vector ops; ops.reserve(block.getOperations().size()); diff --git a/tensorflow/compiler/mlir/xla/transforms/lhlo_fuse_linalg.cc b/tensorflow/compiler/mlir/xla/transforms/lhlo_fuse_linalg.cc index e16ab571b4d..f0971fdf76e 100644 --- a/tensorflow/compiler/mlir/xla/transforms/lhlo_fuse_linalg.cc +++ b/tensorflow/compiler/mlir/xla/transforms/lhlo_fuse_linalg.cc @@ -19,6 +19,7 @@ limitations under the License. #include "mlir/Dialect/Linalg/Analysis/DependenceAnalysis.h" #include "absl/memory/memory.h" #include "llvm/ADT/ArrayRef.h" +#include "llvm/ADT/STLExtras.h" #include "mlir/Dialect/Linalg/Transforms/Transforms.h" // from @llvm-project #include "mlir/Dialect/StandardOps/IR/Ops.h" // from @llvm-project #include "mlir/Pass/Pass.h" // from @llvm-project @@ -44,7 +45,7 @@ class LhloFuseLinalg : public PassWrapper { auto func = getFunction(); // TODO(pifon): Remove assumption that the function has a single block. - if (func.getBlocks().size() != 1) { + if (!llvm::hasSingleElement(func)) { emitError(func.getLoc(), "The function needs to have a single block."); signalPassFailure(); return; @@ -58,7 +59,7 @@ class LhloFuseLinalg : public PassWrapper { for (auto func_arg : func.getArguments()) { result_buffers.insert(func_arg); } - for (auto& block : func.getBlocks()) { + for (auto& block : func) { auto returnOp = mlir::dyn_cast(block.getTerminator()); if (!returnOp) continue; for (auto operand : returnOp.getOperands()) { diff --git a/tensorflow/compiler/mlir/xla/transforms/xla_hlo_fusion.cc b/tensorflow/compiler/mlir/xla/transforms/xla_hlo_fusion.cc index c4eb0e143d2..5d3eda0bea5 100644 --- a/tensorflow/compiler/mlir/xla/transforms/xla_hlo_fusion.cc +++ b/tensorflow/compiler/mlir/xla/transforms/xla_hlo_fusion.cc @@ -487,7 +487,7 @@ struct XlaHloFusion : public mlir::PassWrapper { } // process each block and do fusion within a block. - for (Block& block : func.getBlocks()) { + for (Block& block : func) { SmallVector op_list; for (Operation& op : block) { op_list.push_back(&op); diff --git a/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc b/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc index ecd1308be4b..3f99d40c717 100644 --- a/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc +++ b/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc @@ -301,7 +301,7 @@ struct RewriteKernelSignature signalPassFailure(); return; } - if (func.getBlocks().size() != 1) { + if (!llvm::hasSingleElement(func)) { func.emitError() << "surrounding function has more than one block"; signalPassFailure(); return;