namespace tensorflow {
namespace jit {

void DeviceSet::Insert(DeviceId device_id) {
  int word_index = device_id.id() / kWordSize;
  int bit_index = device_id.id() % kWordSize;
  const int storage_size = storage_.size();
  if (word_index >= storage_size) {
    storage_.resize(word_index + 1, 0);
  }

  storage_[word_index] |= (1ull << bit_index);
}

void DeviceSet::UnionWith(const DeviceSet& other) {

  explicit DeviceId(int id) : id_(id) {}

  int id() const { return id_; }

  friend class DeviceInfoCache;
  friend class DeviceSet;
};

// A set of DeviceIds, represented as a bitmap.
class DeviceSet {
 public:
  void Insert(DeviceId device_id);
  void UnionWith(const DeviceSet& other);
  bool IsEmpty() const;

  // Calls `func` on each DeviceId in the set.  Stops iterating early if `func`
  // return false.

class MockMlirOptimizationPass : public MlirOptimizationPass {
 public:
  MOCK_METHOD(llvm::StringRef, name, (), (const, override));
  MOCK_METHOD(MlirOptimizationPassState, GetPassState,
              (const DeviceSet* device_set, const ConfigProto& config_proto,
               const Graph& graph,
               const FunctionLibraryDefinition& function_library),
              (const, override));
  MOCK_METHOD(Status, Run,

namespace {

Status PickDeviceHelper(bool allow_mixing_unknown_and_cpu,
                        absl::Span<const absl::string_view> device_names,
                        string* result) {
  jit::DeviceInfoCache cache;
  jit::DeviceSet device_set;
  for (absl::string_view name : device_names) {
    TF_ASSIGN_OR_RETURN(jit::DeviceId device_id, cache.GetIdFor(name));
    device_set.Insert(device_id);
  }

  TF_ASSIGN_OR_RETURN(

    } else {
      devices->AddDevice(device);
    }
  }

  return mlir::success();
}

}  // namespace

void AddDevicesToOp(mlir::Operation* op, const DeviceSet* device_set) {
  if (!device_set) return;

  mlir::MLIRContext* ctx = op->getContext();
  mlir::Builder builder(ctx);

  // Collect devices with attached metadata.

                                "The number of graphs being op expanded.");
}  // namespace

namespace tfr {

MlirOptimizationPassState GraphDecomposePass::GetPassState(
    const DeviceSet* device_set, const ConfigProto& config_proto,
    const Graph& graph,
    const FunctionLibraryDefinition& function_library) const {
  const char* tfr_lib_env_val = getenv(std::string(kTFRLibEnv).c_str());

#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/util/device_name_utils.h"

namespace tensorflow {
namespace {

// A fake device used to populate a DeviceSet.
class FakeDevice : public Device {
 public:
  explicit FakeDevice(const DeviceAttributes& device_attributes)
      : Device(nullptr, device_attributes) {}

  std::vector<std::unique_ptr<Device>> devices;
  devices.push_back(FakeDevice::Make(kDeviceName, DEVICE_GPU));
  devices.push_back(FakeDevice::Make(kHostName, DEVICE_CPU));

  std::unique_ptr<DeviceSet> device_set(new DeviceSet());
  for (auto& device : devices) {
    device_set->AddDevice(device.get());
  }

  auto graph = std::make_unique<Graph>(OpRegistry::Global());
  SessionOptions session_options;

#include "tensorflow/core/public/version.h"
#include "tensorflow/core/util/dump_graph.h"

namespace tensorflow {

namespace {
using DeadnessPredicate = DeadnessAnalysis::DeadnessPredicate;
using jit::DeviceId;
using jit::DeviceSet;

// The clusters we create here are eventually lowered into an
// _XlaCompile/_XlaRun pair with a TF executor "fallback" that uses the
// PartitionedCall op to execute the cluster in the regular graph executor if

  mlir::func::registerAllExtensions(registry);
  // clang-format on
}

Status MlirFunctionOptimizationPass::Run(
    const std::string& function_name, const DeviceSet& device_set,
    const ConfigProto& config_proto,
    const FunctionOptimizationPass::FunctionOptions& function_options,
    std::unique_ptr<Graph>* graph, FunctionLibraryDefinition* flib_def,