#include "tensorflow/core/common_runtime/optimization_registry.h"

namespace tensorflow {

// An optimization pass which marks the constants which have to be resolved for
// XLA compilation with `_input_hostmem`.
class ForceXlaConstantsOnHostPass : public GraphOptimizationPass {
 public:
  ForceXlaConstantsOnHostPass() = default;

  Status Run(const GraphOptimizationPassOptions& options) override;
};

          flr, function, &fbody, &constant_arg_indices, &resource_arg_indices));
      VLOG(3) << "Found constant arg indices: "
              << absl::StrJoin(constant_arg_indices, ", ");

      node->AddAttr("_input_hostmem", constant_arg_indices);
    }
  }
  return absl::OkStatus();
}

  for (Node* node : graph->nodes()) {
    if (CanCreateXlaKernel(node->def())) {
      EXPECT_FALSE(found);
      found = true;
      std::vector<int32> hostmem_attr;
      EXPECT_TRUE(TryGetNodeAttr(node->def(), "_input_hostmem", &hostmem_attr));
      EXPECT_EQ(hostmem_attr.size(), 1);
      EXPECT_EQ(hostmem_attr[0], 1);
    }
  }
  EXPECT_TRUE(found);
}

}  // namespace