// RUN: tf-tfrt-opt %s -tf-restore-splitting | FileCheck %s

// CHECK-LABEL: func @single_restore
// CHECK-SAME:    (%[[ARG0:.*]]: {{.*}})
func.func @single_restore(%arg0: tensor<!tf_type.string>) -> (tensor<*xf32>, tensor<*xi32>) {
  %0 = "tf.Const"() {value = dense<["foo", "bar"]> : tensor<2x!tf_type.string>} : () -> tensor<2x!tf_type.string>
  %1 = "tf.Const"() {value = dense<""> : tensor<2x!tf_type.string>} : () -> tensor<2x!tf_type.string>

    ::testing::Combine(::testing::Values("graphdef", "mlir"),
                       /*tfrt*/ ::testing::Values(false),
                       /*use_function*/ ::testing::Values(true, false)));
#else
INSTANTIATE_TEST_SUITE_P(
    UnifiedCAPI, CppGradients,
    ::testing::Combine(::testing::Values("graphdef", "mlir"),
                       /*tfrt*/ ::testing::Values(false),

// RUN: tf-tfrt-opt -tfrt-deduplicate-functions-invoked-by-batch-function %s | FileCheck %s

// This test verifies the function `compute_1` will be removed to deduplicate
// the functions invoked by BatchFunction with the same shared_name and the
// function `compute_2` will not be removed as the shared_name is different.

// CHECK-LABEL: func private @batch_0
// CHECK: f = @compute_0
func.func private @batch_0(%arg0: tensor<?x?xi32>) -> tensor<*xi32> {

See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/tfrt/analysis/tensor_array_side_effect_analysis.h"

#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
#include "mlir/Interfaces/SideEffectInterfaces.h"  // from @llvm-project

#include "tensorflow/core/platform/status_matchers.h"
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/core/protobuf/error_codes.pb.h"
#include "tensorflow/core/tfrt/common/create_pjrt_client_util.h"
#include "tensorflow/core/tfrt/common/pjrt_util.h"
#include "tensorflow/core/tpu/tpu_defs.h"

namespace tensorflow {
namespace {
using XlaDeviceCompiler =
    DeviceCompiler<xla::LocalExecutable, xla::LocalClient>;

// RUN: tf-tfrt-opt %s -tf-identity-propagation -canonicalize | FileCheck %s

// CHECK-LABEL: func @identity
// CHECK-SAME:    (%[[ARG0:.*]]: tensor<i32>)
func.func @identity(%arg0: tensor<i32>) -> tensor<i32> {
  // CHECK-NOT: "tf.Identity"
  %0 = "tf.Identity"(%arg0) : (tensor<i32>) -> tensor<i32>
  // CHECK: return %[[ARG0]]
  func.return %0 : tensor<i32>
}

// CHECK-LABEL: func @identity_terminator
// CHECK-SAME:    (%[[ARG0:.*]]: tensor<i32>)

#define TENSORFLOW_COMPILER_MLIR_TFRT_BACKEND_COMPILER_H_

#include "absl/status/status.h"
#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
#include "mlir/IR/DialectRegistry.h"  // from @llvm-project
#include "tensorflow/core/tfrt/runtime/runtime.h"

namespace tensorflow {

class BackendCompiler {
 public:
  virtual ~BackendCompiler();

  virtual void GetDependentDialects(mlir::DialectRegistry& registry) const {}

#include "xla/pjrt/pjrt_api.h"
#include "xla/pjrt/pjrt_c_api_client.h"
#include "xla/shape.h"
#include "xla/shape_util.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/tfrt/common/async_value_tensor.h"
#include "tensorflow/core/tfrt/common/pjrt_util.h"
#include "tsl/lib/core/status_test_util.h"
#include "tsl/platform/casts.h"
#include "tsl/platform/status_matchers.h"
#include "tsl/protobuf/error_codes.pb.h"

// RUN: tf-tfrt-opt -split-input-file -tf-mlrt-rewrite-ifrt-load-variable %s | FileCheck %s

// Variable is used by both CPU and TPU
//
// CHECK-LABEL: func @serving_default(%arg0: tensor<1x3xf32>) -> tensor<1x1xf32>
// CHECK-NEXT:    [[HANDLE:%.*]] = "tf.VarHandleOp"()
// CHECK-NEXT:    [[ARRAYKEY:%.*]], [[FURTURE:%.*]] = "tf_mlrt.tf_ifrt_load_variable"([[HANDLE]])

  static llvm::StringRef getDialectNamespace() { return "tf_mlrt_tpu"; }
};

}  // namespace tf_mlrt_tpu
}  // namespace tensorflow

#define GET_OP_CLASSES
#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_tpu_ops.h.inc"