using mlir::WalkResult;
using mlir::func::FuncOp;

constexpr char kXlaMapOutsideCompilationAttr[] = "_xla_map_outside_compilation";
constexpr char kXlaOutsideCompilationAttr[] = "_xla_outside_compilation";

// Return true if `op` has attributes that say it can be outside compiled by
// this pass. This pass ignores _xla_map_outside_compilation, which will only be
// handled by extract_outside_compilation pass.

    // CHECK:             "tf.OpC"(%[[B_FULL]])
    "tf_device.cluster"() ({
      %0 = "tf.OpA"() {_XlaSharding = "\08\03\1A\02\02\01\22\02\00\01"} : () -> tensor<2x2xi64>
      %1 = "tf.OpB"(%0) {_xla_map_outside_compilation = "0", _xla_outside_compilation = "from_launch"} : (tensor<2x2xi64>) -> tensor<2x2xi64>
      "tf.OpC"(%1) : (tensor<2x2xi64>) -> ()
      tf_device.return

using mlir::func::FuncOp;
using mlir::func::ReturnOp;

constexpr char kDeviceAttr[] = "device";
constexpr char kHostFunctionAttr[] = "host_func";
constexpr char kXlaMapOutsideCompilationAttr[] = "_xla_map_outside_compilation";
constexpr char kXlaOutsideCompilationAttr[] = "_xla_outside_compilation";

#define GEN_PASS_DEF_EXTRACTOUTSIDECOMPILATIONPASS
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.h.inc"