Migrate constant sinking pass to use declarative pass registration instead of manually defined pass registration (NFC).

Pass documentation is also improved and migrated to the declarative pass spec. Pass is also renamed to better reflect what ops it operates on (tf_device.cluster). The pass does not operate on the tf_executor dialect.

PiperOrigin-RevId: 348120524
Change-Id: Ic828b720a8a087548c162a1299aef1655b479cd6

Author: andyly

tensorflow/compiler/mlir/tensorflow/transforms/bridge.cc

@@ -110,7 +110,7 @@ void CreateTPUBridgePipeline(OpPassManager &pm) {
   pm.addPass(CreateTPUOutsideCompilationClusterPass());
   pm.addPass(CreateTPUExtractOutsideCompilationPass());
 
-  pm.addNestedPass<FuncOp>(tf_executor::CreateTFExecutorConstantSinkingPass());
+  pm.addNestedPass<FuncOp>(TFDevice::CreateClusterConstantSinkingPass());
   pm.addPass(TF::CreateResourceDeviceInferencePass());
   pm.addPass(TFDevice::CreateClusterOutliningPass());
   pm.addPass(CreateTPUDynamicPaddingMapperPass());

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

@@ -237,18 +237,18 @@ CreateTFExecutorTPUV1IslandInliningPass();
 
 // Creates a pass to prune tf_executor.graph from dead nodes.
 std::unique_ptr<OperationPass<FuncOp>> CreateTFExecutorGraphPruningPass();
-
-// Sink `tf.Const` operations in the LaunchOp region using them. This is
-// performed in order to limit the number of values implicitly captured in this
-// region before outlining.
-std::unique_ptr<OperationPass<FuncOp>> CreateTFExecutorConstantSinkingPass();
 }  // namespace tf_executor
 
 namespace TFDevice {
 // Creates a pass that forms clusters from instructions that are assigned to
 // same device.
 std::unique_ptr<OperationPass<FuncOp>> CreateClusterFormationPass();
 
+// Sinks `tf.Const` operations in the ClusterOp region using them. This is
+// performed in order to limit the number of values implicitly captured in this
+// region before outlining.
+std::unique_ptr<OperationPass<FuncOp>> CreateClusterConstantSinkingPass();
+
 // Creates a pass that outlines regions of tf_device.launch operations.
 std::unique_ptr<OperationPass<ModuleOp>> CreateClusterOutliningPass();
 

tensorflow/compiler/mlir/tensorflow/transforms/sink_constant.cc

@@ -25,21 +25,21 @@ limitations under the License.
 #include "mlir/Transforms/Passes.h"  // from @llvm-project
 #include "mlir/Transforms/RegionUtils.h"  // from @llvm-project
 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"
-#include "tensorflow/compiler/mlir/tensorflow/ir/tf_executor.h"
 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
 #include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
+#include "tensorflow/compiler/mlir/tensorflow/transforms/passes_detail.h"
 #include "tensorflow/compiler/mlir/tensorflow/utils/error_util.h"
 
 #define DEBUG_TYPE "tf-executor-sink-constant"
 
 namespace mlir {
-namespace tf_executor {
+namespace TFDevice {
 
 namespace {
 using ::mlir::TF::ConstOp;
 
-class ExecutorConstantSinking
-    : public mlir::PassWrapper<ExecutorConstantSinking, FunctionPass> {
+class ClusterConstantSinkingPass
+    : public TF::ClusterConstantSinkingPassBase<ClusterConstantSinkingPass> {
   void runOnFunction() override {
     getFunction().walk([](tf_device::ClusterOp cluster) {
       LLVM_DEBUG(llvm::dbgs() << "Visit " << *cluster.getOperation() << "\n");
@@ -82,16 +82,11 @@ class ExecutorConstantSinking
   }
 };
 
-static mlir::PassRegistration<ExecutorConstantSinking> pass(
-    "tf-device-constant-sinking",
-    "Sink constants implicitly captured in a tf_device.cluster region. This "
-    "reduces the number of arguments when outlining later.");
-
 }  // anonymous namespace
 
-std::unique_ptr<OperationPass<FuncOp>> CreateTFExecutorConstantSinkingPass() {
-  return std::make_unique<ExecutorConstantSinking>();
+std::unique_ptr<OperationPass<FuncOp>> CreateClusterConstantSinkingPass() {
+  return std::make_unique<ClusterConstantSinkingPass>();
 }
 
-}  // namespace tf_executor
+}  // namespace TFDevice
 }  // namespace mlir

tensorflow/compiler/mlir/tensorflow/transforms/tf_passes.td

@@ -202,6 +202,44 @@ func @tpu_computation(%arg0: tensor<i32>, %arg1: tensor<i32>) -> (tensor<i32>, t
   let constructor = "TFTPU::CreateTPUClusterFormationPass()";
 }
 
+def ClusterConstantSinkingPass : FunctionPass<"tf-device-constant-sinking"> {
+  let summary = "Sinks constants implicitly captured in a tf_device.cluster region.";
+
+  let description = [{
+This pass sinks implicitly captured constants (`tf.Const` ops) used by and into
+a `tf_device.cluster` region. Performing this prior to outlining will reduce the
+number of arguments of the outlined function.
+
+For example, the following:
+
+```mlir
+func @cluster() -> tensor<i32> {
+  %const = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
+  %cluster = "tf_device.cluster"() ( {
+    %identity = "tf.Identity"(%const) : (tensor<i32>) -> tensor<i32>
+    tf_device.return %identity : tensor<i32>
+  }) : () -> (tensor<i32>)
+  return %cluster : tensor<i32>
+}
+```
+
+will be transformed into:
+
+```mlir
+func @cluster() -> tensor<i32> {
+  %cluster = "tf_device.cluster"() ( {
+    %const = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
+    %identity = "tf.Identity"(%const) : (tensor<i32>) -> tensor<i32>
+    tf_device.return %identity : tensor<i32>
+  }) : () -> (tensor<i32>)
+  return %cluster : tensor<i32>
+}
+```
+  }];
+
+  let constructor = "TFDevice::CreateClusterConstantSinkingPass()";
+}
+
 def TPUExtractOutsideCompilationPass : Pass<"tf-tpu-extract-outside-compilation", "ModuleOp"> {
   let summary = "Extracts TPU outside compilation computation to a separate tf_device.parallel_execute region.";