namespace TF {

std::string GetDeviceAttrAsResourceInstanceStr(Operation* op);

void MarkResourceAsReadAndWrite(
    Value value,
    SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>&
        effect);

void MarkResourceAsReadOnly(
    Value value,
    SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>&
        effect);

}  // namespace TF
}  // namespace mlir

  // string) are seen as dependent (same resource instance).
  if (!device_attr) return "";
  return device_attr.str();
}

void MarkResourceAsReadAndWrite(
    Value value,
    SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>&
        effects) {
  if (value.getType().cast<TensorType>().getElementType().isa<ResourceType>()) {
    effects.emplace_back(MemoryEffects::Read::get(), value,

                                  SmallVectorImpl<Type>* output_types) const;

  // Creates the TF op and also the necessary tfr.cast ops to replace the
  // original TFR call op.
  LogicalResult CreateAndReplaceOp(
      PatternRewriter& rewriter, CallOp call_op,
      const SmallVectorImpl<Type>& output_types,
      const SmallVectorImpl<Value>& inputs, const NamedAttrList& attr_list,

                                 llvm::SmallVectorImpl<Type>& new_input_types,
                                 OpBuilder builder);

  // Modifies the element types of entry block returns to be user specified
  // and returns the new return types.
  LogicalResult ModifyOutputNodes(func::FuncOp func,
                                  llvm::SmallVectorImpl<Type>& new_output_types,

void ConvertToMlirShape(const TensorShape& input_shape,
                        llvm::SmallVectorImpl<int64_t>* shape);

// Converts an TensorFlow shape proto to the one used in MLIR.
Status ConvertToMlirShape(const TensorShapeProto& input_shape,
                          llvm::SmallVectorImpl<int64_t>* shape);

// Given a tensor shape and dtype, get the corresponding MLIR tensor type.

mlir::LogicalResult GetOutputTypesForLogicalDeviceComputation(
    int core_id, llvm::ArrayRef<xla::OpSharding> output_sharding_config,
    mlir::tf_device::ClusterFuncOp cluster_func,
    llvm::SmallVectorImpl<mlir::Type>* output_types,
    llvm::SmallVectorImpl<int>* cluster_to_core_index);

// Remaps outputs of `new_parallel_execute` op that represent concurrent
// execution of the `tf_device.cluster_func` at index `cluster_idx` of

  LogicalResult RewriteOp(
      Operation* op, PatternRewriter& rewriter,
      const std::function<bool(ArrayRef<int64_t>, ArrayRef<int64_t>,
                               SmallVectorImpl<int64_t>&)>&
          get_broadcasted_shape) const;

  LogicalResult RewriteBatchMatMulV2Op(Operation* op,
                                       PatternRewriter& rewriter) const;
};

                                             SmallVectorImpl<OpFoldResult> &);
  static void RegisterConstantFoldHook(ConstantFoldHook fn) {
    constant_fold_hook_ = std::move(fn);
  }

  static LogicalResult constantFold(Operation *op, ArrayRef<Attribute> operands,
                                    SmallVectorImpl<OpFoldResult> &results) {

// tf_executor.graph.
void VisitOpOperand(GraphOp graph, Value operand,
                    llvm::SmallPtrSetImpl<Operation*>* reachable_ops,
                    llvm::SmallVectorImpl<Operation*>* ops_to_visit) {
  Operation* def = operand.getDefiningOp();
  if (def && def->getParentOp() == graph && reachable_ops->insert(def).second) {
    // Op has not been visited, add to queue to visit later.

#include "mlir/Support/LogicalResult.h"  // from @llvm-project

namespace mlir {
namespace TF {

LogicalResult ConstantFoldFallbackHook(
    Operation *inst, ArrayRef<Attribute> operands,
    SmallVectorImpl<OpFoldResult> &results);  // NOLINT

}  // namespace TF
}  // namespace mlir