#include "tensorflow/compiler/jit/defs.h"

#include <atomic>

namespace tensorflow {

const char* const kXlaMustCompileAttr = "_XlaMustCompile";

const char* const kXlaCompileAttr = "_XlaCompile";

// User-provided through jit_scope APIs. Effective only when auto_jit is OFF.
const char* const kXlaScopeAttr = "_XlaScope";

// Automatically inserted by auto_jit to guide clustering results.  Effective

extern const char* const kXlaMustCompileAttr;  // "_XlaMustCompile"

// Implies auto-clustering: tagged nodes will be clustered and compiled with XLA
// on a best-effort basis.
extern const char* const kXlaCompileAttr;  // "_XlaCompile"

// Implies auto-clustering within the given scope.
extern const char* const kXlaScopeAttr;    // "_XlaScope"
extern const char* const kXlaInternalScopeAttr;  // "_XlaInternalScope"

 public:
  // Given a NodeDef 'node_def' and the function library runtime 'flr', returns
  // true if 'node_def' is a call to a compilable function defined in 'flr',
  // with the kXlaCompileAttr set.
  bool CanCreateKernel(
      const FunctionLibraryRuntime& flr,
      const std::shared_ptr<const NodeProperties>& props) const override;

                                         .WithAttr(kXlaCompileAttr, true)
                                         .WithAttr(kXlaScopeAttr, "Scope1"));
    Node* b = ops::UnaryOp("Relu", a,
                           builder.opts()
                               .WithName("B")
                               .WithAttr(kXlaCompileAttr, true)
                               .WithAttr(kXlaScopeAttr, "Scope1"));

  bool compile = false;
  Status status = GetNodeAttr(node->attrs(), kXlaCompileAttr, &compile);
  if (status.ok()) {
    if (!compile) {
      VLOG(2) << "Rejecting " << node->name() << ": kXlaCompileAttr("
              << kXlaCompileAttr << ") is false.";
    }
    return !compile;
  }

  status = flib_def_->GetAttr(*node, kXlaCompileAttr, &compile);
  if (status.ok()) {
    if (!compile) {

                                       std::vector<int>* constant_arg_indices,
                                       std::vector<int>* resource_arg_indices);

// Given a NodeDef `node_def` returns true iff `node_def` has kXlaCompileAttr
// set.
bool CanCreateXlaKernel(const NodeDef& node_def);

// Returns memory types for the input.
// `constant_arg_indices` and `resource_arg_indices` are sorted arrays of

static bool NodeCanTriggerXlaCompilation(const NodeDef& node) {
  return node.attr().find(kXlaClusterIdAttr) != node.attr().end() ||
         HasBoolAttr(node, kXlaMustCompileAttr) ||
         HasBoolAttr(node, kXlaCompileAttr) ||
         HasBoolAttr(node, kXlaScopeAttr) ||
         HasBoolAttr(node, kXlaInternalScopeAttr) ||
         ops_triggering_xla_compilation->count(node.op());
}