Migrate control flow conversion passes to TableGen

Now FunctionalControlFlowToRegions and RegionControlFlowToFunctional passes
are defined and documented in TableGen.

PiperOrigin-RevId: 348863824
Change-Id: I94dfabd2582dc0e00252d808794ee19f8ae87932

Author: mgester

tensorflow/compiler/mlir/tensorflow/transforms/functional_control_flow_to_regions.cc

@@ -33,6 +33,7 @@ limitations under the License.
 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"
 #include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
+#include "tensorflow/compiler/mlir/tensorflow/transforms/passes_detail.h"
 #include "tensorflow/compiler/mlir/tensorflow/utils/attribute_utils.h"
 
 #define DEBUG_TYPE "tf-functional-cf-to-region"
@@ -43,8 +44,8 @@ namespace TF {
 namespace {
 
 struct FunctionalControlFlowToRegions
-    : public PassWrapper<FunctionalControlFlowToRegions,
-                         OperationPass<ModuleOp>> {
+    : public TF::FunctionalControlFlowToRegionsPassBase<
+          FunctionalControlFlowToRegions> {
   void runOnOperation() override;
 };
 
@@ -157,9 +158,5 @@ CreateTFFunctionalControlFlowToRegions() {
   return std::make_unique<FunctionalControlFlowToRegions>();
 }
 
-static PassRegistration<FunctionalControlFlowToRegions> pass(
-    "tf-functional-control-flow-to-regions",
-    "Transform functional control flow Ops to Region based counterparts");
-
 }  // namespace TF
 }  // namespace mlir

tensorflow/compiler/mlir/tensorflow/transforms/region_control_flow_to_functional.cc

@@ -36,6 +36,7 @@ limitations under the License.
 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"
 #include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
+#include "tensorflow/compiler/mlir/tensorflow/transforms/passes_detail.h"
 #include "tensorflow/compiler/mlir/tensorflow/utils/attribute_utils.h"
 
 #define DEBUG_TYPE "tf-region-cf-to-functional"
@@ -46,8 +47,8 @@ namespace TF {
 namespace {
 
 struct RegionControlFlowToFunctional
-    : public PassWrapper<RegionControlFlowToFunctional,
-                         OperationPass<ModuleOp>> {
+    : public TF::RegionControlFlowToFunctionalPassBase<
+          RegionControlFlowToFunctional> {
   void runOnOperation() override;
 
  private:
@@ -445,9 +446,5 @@ CreateTFRegionControlFlowToFunctional() {
   return std::make_unique<RegionControlFlowToFunctional>();
 }
 
-static PassRegistration<RegionControlFlowToFunctional> pass(
-    "tf-region-control-flow-to-functional",
-    "Transform region bases control flow Ops to functional counterparts");
-
 }  // namespace TF
 }  // namespace mlir

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

@@ -427,3 +427,67 @@ func @unsupported_op() -> tensor<i32> {
 
   let constructor = "TFDevice::CreateMarkOpsForOutsideCompilationPass()";
 }
+
+def FunctionalControlFlowToRegionsPass : Pass<"tf-functional-control-flow-to-regions", "ModuleOp"> {
+  let summary = "Transforms functional control flow operations to their region-based counterparts";
+
+  let description = [{
+This pass transforms functional control flow operations in the TensorFlow
+dialect to their region-based counterparts, i.e., `tf.If` is transformed to
+`tf.IfRegion` and `tf.While` is transformed to `tf.WhileRegion`.
+
+For example, this functional operation
+
+```mlir
+  %0 = "tf.If"(%arg0, %arg1) {
+    then_branch = @then_branch_func, else_branch = @else_branch_func, is_stateless = false
+  } : (tensor<i1>, tensor<*xf32>) -> tensor<*xf32>
+```
+
+will be transformed into this region-based operation
+
+```mlir
+    %0 = "tf.IfRegion"(%arg0) ( {
+      %1 = call @then_branch_func(%arg1) : (tensor<*xf32>) -> tensor<*xf32>
+      "tf.Yield"(%1) : (tensor<*xf32>) -> ()
+    },  {
+      %1 = call @else_branch_func(%arg1) : (tensor<*xf32>) -> tensor<*xf32>
+      "tf.Yield"(%1) : (tensor<*xf32>) -> ()
+    }) {is_stateless = false} : (tensor<i1>) -> tensor<*xf32>
+```
+  }];
+
+  let constructor = "TF::CreateTFFunctionalControlFlowToRegions()";
+}
+
+def RegionControlFlowToFunctionalPass : Pass<"tf-region-control-flow-to-functional", "ModuleOp"> {
+  let summary = "Transforms region-based control flow operations to their functional counterparts";
+
+  let description = [{
+This pass transforms region-based control flow operations in the TensorFlow
+dialect to their functional counterparts, i.e., `tf.IfRegion` is transformed to
+`tf.If` and `tf.WhileRegion` is transformed to `tf.While`.
+
+For example, this region-based operation
+
+```mlir
+    %0 = "tf.IfRegion"(%arg0) ( {
+      %1 = call @then_branch_func(%arg1) : (tensor<*xf32>) -> tensor<*xf32>
+      "tf.Yield"(%1) : (tensor<*xf32>) -> ()
+    },  {
+      %1 = call @else_branch_func(%arg1) : (tensor<*xf32>) -> tensor<*xf32>
+      "tf.Yield"(%1) : (tensor<*xf32>) -> ()
+    }) {is_stateless = false} : (tensor<i1>) -> tensor<*xf32>
+```
+
+will be transformed into this functional operation
+
+```mlir
+  %0 = "tf.If"(%arg0, %arg1) {
+    then_branch = @then_branch_func, else_branch = @else_branch_func, is_stateless = false
+  } : (tensor<i1>, tensor<*xf32>) -> tensor<*xf32>
+```
+  }];
+
+  let constructor = "TF::CreateTFRegionControlFlowToFunctional()";
+}