#include "tensorflow/compiler/aot/compile.h"
#include "tensorflow/compiler/aot/flags.h"
#include "tensorflow/compiler/tf2xla/tf2xla.pb.h"
#include "tensorflow/compiler/tf2xla/tf2xla_util.h"
#include "xla/debug_options_flags.h"
#include "xla/service/compiler.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/tensor_shape.h"

syntax = "proto3";

package tensorflow;

import "xla/service/hlo.proto";

// Represents the cache key used for persistence.
message XlaSerializedCacheKey {
  uint64 signature_fingerprint = 1;
  uint64 cluster_fingerprint = 2;
  string device_type = 3;
  string prefix = 4;
  bool compiled_using_pjrt = 5;
}

// Represents an entry in the XLA compile cache.
message XlaSerializedCacheEntry {

// Checks if a device string is a TPU replicated core device.
bool IsTPUReplicatedCore(llvm::StringRef device);

// Checks if `type` is allowed for XLA. String and resources are not XLA types.
// There are other TF types that are not XLA types which will be removed by
// successive passes in TF/XLA bridge phase 2.
bool TypeValidForXLA(const mlir::Type& type);

// Returns the map from core to the host that is associated with the

# MLIR dialects and utilities for TensorFlow, TensorFlow Lite and XLA.

This module contains the MLIR
([Multi-Level Intermediate Representation](https://mlir.llvm.org))
dialects and utilities for

1. TensorFlow
2. XLA
3. TF Lite

See [MLIR's website](https://mlir.llvm.org) for complete documentation.

## Getting started

Building dialects and utilities here follow the standard approach using
`bazel` as the rest of TensorFlow.

}

def XlaInlineDeviceOpsPass : Pass<"tf-xla-inline-device-ops", "ModuleOp"> {
  let summary = "Inline all Cluster op based in the parent region";
  let constructor = "TFDevice::CreateXlaInlineDeviceOpsPass()";
  let dependentDialects = ["tf_device::TensorFlowDeviceDialect"];
}

def XlaRewritePass : Pass<"tf-xla-rewrite", "mlir::ModuleOp"> {
  let summary = "Rewrites partition calls into Xla launch ops to make the attached function run on XLA.";

#include "tensorflow/compiler/mlir/tf2xla/transforms/passes.h"
#include "tensorflow/compiler/mlir/tf2xla/transforms/xla_legalize_targets.h"
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
#include "xla/mlir_hlo/mhlo/transforms/rewriters.h"
#include "xla/mlir_hlo/mhlo/utils/type_conversion.h"
#include "tensorflow/core/lib/monitoring/counter.h"

namespace mlir {
namespace mhlo {
namespace {

  // program_shape_access_shim is a C++ expression that constructs the
  // xla::ProgramShapeProto instance for the CompileResult passed to
  // GenerateMetadata.
  string program_shape_access_shim;

  // hlo_profile_printer_data_access_shim is a C++ expression that constructs
  // the xla::HloProfilePrinterData instance for the CompileResult passed to
  // GenerateMetadata.  If the xla::HloProfilePrinterData is null then this is a

#include "llvm/ADT/StringRef.h"
#include "mlir/Pass/Pass.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf.h"
#include "xla/client/client_library.h"
#include "xla/shape.h"
#include "xla/stream_executor/platform.h"
#include "xla/stream_executor/platform_manager.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/types.pb.h"

#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/status/status.h"
#include "tensorflow/compiler/mlir/tf2xla/api/v1/compile_mlir_util.h"
#include "tensorflow/compiler/tf2xla/xla_helpers.h"
#include "xla/client/xla_computation.h"
#include "xla/service/hlo.pb.h"
#include "tensorflow/core/lib/monitoring/cell_reader.h"
#include "tensorflow/core/lib/monitoring/counter.h"
#include "tsl/platform/statusor.h"

namespace {

// RUN: stablehlo-quant-opt %s -split-input-file -tf-xla-call-module-serialization -stablehlo-convert-xla-call-module-op-to-bfloat16 -tf-xla-call-module-deserialization | FileCheck %s

// ConvertXlaCallModuleOpToBfloat16Pass works on XlaCallModuleOps with
// serialized modules. Which makes verification difficult. Therefore we add
// (de)serialization passes so that the input and output are deserializated
// StableHLO functions.

// CHECK-LABEL: module
module {