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,

  // `function_library` contains function definitions for function calls in
  // `graph` not included in the `graph` FunctionLibraryDefinition.
  virtual MlirOptimizationPassState GetPassState(
      const DeviceSet* device_set, const ConfigProto& config_proto,
      const Graph& graph,
      const FunctionLibraryDefinition& function_library) const = 0;

  virtual Status Run(const std::string& function_name,

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.

 public:
  llvm::StringRef name() const override { return "graph_optimization"; }

  ::tensorflow::MlirOptimizationPassState GetPassState(
      const ::tensorflow::DeviceSet* device_set,
      const ::tensorflow::ConfigProto& config_proto,
      const tensorflow::Graph& graph,
      const tensorflow::FunctionLibraryDefinition& function_library)
      const override {

  // Whether to run this pass. If this is enabled, the GraphDef will be imported
  // to MLIR even no tf composition file is found.
  ::tensorflow::MlirOptimizationPassState GetPassState(
      const DeviceSet* device_set, const ConfigProto& config_proto,
      const Graph& graph,
      const FunctionLibraryDefinition& function_library) const override;

  // This should be used as a thin mapper around mlir::ModulePass::runOnModule

//   /job:<name>/replica:<replica>/task:<task>/device:<type>:<device_num>
//
// Supported device metadata types:
// (1) GpuDeviceMetadata: GPU device compute capability.
void AddDevicesToOp(mlir::Operation* op, const DeviceSet* device_set);

// Collects devices information from an op `tf.devices` attributes. Returns
// failure if can't parse device metadata from the attribute.
mlir::LogicalResult GetDevicesFromOp(mlir::Operation* op,

                                "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());