Minor refactor: rename the 'lower bound batch threads' transform to a more generic 'reconfig batch op'. It makes no logical changes.

tensorflow/compiler/mlir/tfrt/BUILD

@@ -208,12 +208,12 @@ cc_library(
         "transforms/deduplicate_if_result_pass.cc",
         "transforms/fuse_tpu_compile_and_execute_ops.cc",
         "transforms/insert_tensor_copy.cc",
-        "transforms/lower_bound_batch_threads.cc",
         "transforms/lower_saved_model.cc",
         "transforms/merge_tf_if_ops.cc",
         "transforms/optimize.cc",
         "transforms/optimize_tf_control_flow_side_effect.cc",
         "transforms/passes.cc",
+        "transforms/reconfig_batch_op.cc",
         "transforms/remove_device_attribute.cc",
         "transforms/remove_tf_if_const_args.cc",
         "transforms/reorder_assert.cc",

tensorflow/compiler/mlir/tfrt/tests/lower_bound_batch_threads.mlir → tensorflow/compiler/mlir/tfrt/tests/reconfig_batch_op.mlir

@@ -1,4 +1,4 @@
-// RUN: tf-tfrt-opt -split-input-file -tfrt-lower-bound-batch-threads="tfrt-min-num-batch-threads=2" %s | FileCheck %s --dump-input=always
+// RUN: tf-tfrt-opt -split-input-file -tfrt-reconfig-batch-op="tfrt-min-num-batch-threads=2" %s | FileCheck %s --dump-input=always
 
 // -----
 

tensorflow/compiler/mlir/tfrt/transforms/passes.cc

@@ -118,8 +118,8 @@ void CreateTFExecutorToTFPreInvariantOptimizationPipelineHelper(
   pm.addPass(tfrt_compiler::CreateMergeTfIfOpsPass());
 
   // Lower bound on the number of batch threads in `tf.BatchFunction`.
-  pm.addPass(tfrt_compiler::CreateLowerBoundBatchThreadsPass(
-      options.min_num_batch_threads));
+  pm.addPass(tfrt_compiler::CreateReconfigBatchOpPass(
+      {.min_num_batch_threads = options.min_num_batch_threads}));
 
   // Deduplicate functions invoked by tf.BatchFunction with the same
   // shared_name

tensorflow/compiler/mlir/tfrt/transforms/passes.h

@@ -16,6 +16,7 @@ limitations under the License.
 #ifndef TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_PASSES_H_
 #define TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_PASSES_H_
 
+#include <cstdint>
 #include <memory>
 
 #include "llvm/Support/CommandLine.h"
@@ -67,8 +68,11 @@ std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
 CreateDeduplicateFunctionsInovkedByBatchFunctionPass();
 
 // Create a pass to lower bound the number of threads in tf.BatchFunction.
-std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
-CreateLowerBoundBatchThreadsPass(int64_t min_num_batch_threads);
+struct ReconfigBatchOpPassOptions {
+  int64_t min_num_batch_threads = 1;
+};
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>> CreateReconfigBatchOpPass(
+    ReconfigBatchOpPassOptions options);
 
 // Create a pass to fuse the TPU Ops for TFRT.
 std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>

tensorflow/compiler/mlir/tfrt/transforms/lower_bound_batch_threads.cc → tensorflow/compiler/mlir/tfrt/transforms/reconfig_batch_op.cc

@@ -31,34 +31,31 @@ namespace tensorflow {
 namespace tfrt_compiler {
 namespace {
 
-class LowerBoundBatchThreadsPass
-    : public mlir::PassWrapper<LowerBoundBatchThreadsPass,
+class ReconfigBatchOpPass
+    : public mlir::PassWrapper<ReconfigBatchOpPass,
                                mlir::OperationPass<mlir::ModuleOp>> {
  public:
-  explicit LowerBoundBatchThreadsPass(uint64_t min_num_batch_threads)
-      : mlir::PassWrapper<LowerBoundBatchThreadsPass,
+  explicit ReconfigBatchOpPass(ReconfigBatchOpPassOptions options)
+      : mlir::PassWrapper<ReconfigBatchOpPass,
                           mlir::OperationPass<mlir::ModuleOp>>() {
-    min_num_batch_threads_ = min_num_batch_threads;
+    min_num_batch_threads_ = options.min_num_batch_threads;
   }
-  LowerBoundBatchThreadsPass()
-      : mlir::PassWrapper<LowerBoundBatchThreadsPass,
+  ReconfigBatchOpPass()
+      : mlir::PassWrapper<ReconfigBatchOpPass,
                           mlir::OperationPass<mlir::ModuleOp>>() {}
-  LowerBoundBatchThreadsPass(const LowerBoundBatchThreadsPass& other)
-      : mlir::PassWrapper<LowerBoundBatchThreadsPass,
+  ReconfigBatchOpPass(const ReconfigBatchOpPass& other)
+      : mlir::PassWrapper<ReconfigBatchOpPass,
                           mlir::OperationPass<mlir::ModuleOp>>(other) {}
 
-  LowerBoundBatchThreadsPass& operator=(
-      const LowerBoundBatchThreadsPass& other) = delete;
+  ReconfigBatchOpPass& operator=(const ReconfigBatchOpPass& other) = delete;
 
-  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LowerBoundBatchThreadsPass)
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ReconfigBatchOpPass)
 
  private:
-  llvm::StringRef getArgument() const final {
-    return "tfrt-lower-bound-batch-threads";
-  }
+  llvm::StringRef getArgument() const final { return "tfrt-reconfig-batch-op"; }
 
   llvm::StringRef getDescription() const final {
-    return "Lower bound batch threads for batch ops.";
+    return "Reconfig batch op such as num_batch_threads.";
   }
 
   void runOnOperation() override {
@@ -82,12 +79,12 @@ class LowerBoundBatchThreadsPass
 
 }  // namespace
 
-std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
-CreateLowerBoundBatchThreadsPass(int64_t min_num_batch_threads) {
-  return std::make_unique<LowerBoundBatchThreadsPass>(min_num_batch_threads);
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>> CreateReconfigBatchOpPass(
+    ReconfigBatchOpPassOptions options) {
+  return std::make_unique<ReconfigBatchOpPass>(options);
 }
 
-static mlir::PassRegistration<LowerBoundBatchThreadsPass> register_pass;
+static mlir::PassRegistration<ReconfigBatchOpPass> register_pass;
 
 }  // namespace tfrt_compiler
 }  // namespace tensorflow