// RUN: tac-translate -input-mlir -output-mlir -device-specs=NNAPI %s -o - 2>&1 | FileCheck %s

module {
  // CHECK-LABEL: main
  func.func @main(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32> {
    %0 = "tfl.squared_difference"(%arg0, %arg1) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
    func.return %0 : tensor<4xf32>
    // CHECK:  [[VAL_0:%.*]] = tfl.sub %arg0, %arg1 {fused_activation_function = "NONE"} : tensor<4xf32>

// RUN: tac-opt-all-backends -tfl-device-transform-nnapi %s -split-input-file -verify-diagnostics | FileCheck %s

func.func @mean_4d_keepdim(%arg0: tensor<1x48x48x512xf32>) -> tensor<1x1x1x512xf32> {
  %cst = arith.constant dense<[1, 2]> : tensor<2xi32>
  %0 = "tfl.mean"(%arg0, %cst) {keep_dims = true} : (tensor<1x48x48x512xf32>, tensor<2xi32>) -> tensor<1x1x1x512xf32>
  func.return %0 : tensor<1x1x1x512xf32>
}

#define GEN_PASS_DECL_TPUCLUSTERFORMATIONPASS
#define GEN_PASS_DECL_TPUEXTRACTHEADTAILOUTSIDECOMPILATIONPASS
#define GEN_PASS_DECL_TPUEXTRACTOUTSIDECOMPILATIONPASS
#define GEN_PASS_DECL_TPUSHARDINGIDENTIFICATIONPASS
#define GEN_PASS_DECL_VERIFYCLUSTERINGPASS
#define GEN_PASS_DECL_XLACLUSTERFORMATIONPASS
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.h.inc"

}  // namespace internal
}  // namespace tf2xla

#define GEN_PASS_DECL_LEGALIZETF
#define GEN_PASS_DECL_LEGALIZETFCOLLECTIVE
#define GEN_PASS_DECL_LEGALIZETFMODULEPASS
#define GEN_PASS_DECL_LEGALIZETFTYPESPASS
#define GEN_PASS_DECL_TFXLADEVICESPECIFICTRANSFORMS
#define GEN_PASS_DECL_VERIFYTFXLALEGALIZATION
#include "tensorflow/compiler/mlir/tf2xla/transforms/xla_legalize_tf_passes.h.inc"

#define GEN_PASS_REGISTRATION
#define GEN_PASS_DECL_LEGALIZETFCOMMUNICATIONPASS

                             StringPiece tensor_name, Device* device,
                             Tensor* cpu_tensor, StatusCallback done) override;

  void CopyTensorInSameDevice(const Tensor* input_tensor, Device* device,
                              Tensor* output_tensor,
                              StatusCallback done) const override {
    done(errors::Internal("device->device copy not implemented."));
  }

 private:

  }

  void CopyTensorInSameDevice(const Tensor* input_tensor, Device* device,
                              Tensor* output_tensor,
                              StatusCallback done) const override {
    done(errors::Internal("device->device copy not implemented."));
  }

 private:
  se::Stream* stream_;                        // Not owned.

// after the _TPUCompileMlir op.
std::unique_ptr<OperationPass<func::FuncOp>> CreateEmbeddingProgramKeyPass();

#define GEN_PASS_REGISTRATION
#define GEN_PASS_DECL_EMBEDDINGSEQUENCINGPASS
#define GEN_PASS_DECL_EMBEDDINGPIPELININGPASS
#define GEN_PASS_DECL_EMBEDDINGPROGRAMKEYPASS
#include "tensorflow/compiler/mlir/tensorflow/transforms/sparsecore/sparsecore_passes.h.inc"

}  // namespace TFDevice

  }];
}

// This corresponds to tensorflow::Device* .
def TFDeviceType : DialectType<TensorflowMlrt_Dialect,
    CPred<"$_self.isa<::tensorflow::tf_mlrt::TFDeviceType>()">, "!tf_mlrt.device type">,
    BuildableType<"$_builder.getType<::tensorflow::tf_mlrt::TFDeviceType>()"> {
  let description = [{
    `!tf_mlrt.device type` represents a tensorflow::device.
  }];
}

// Classes and utilities that work with StreamExecutor C API for internal use.
// This includes functions used for device registration and interfaces needed
// for testing.
#ifndef TENSORFLOW_C_EXPERIMENTAL_STREAM_EXECUTOR_STREAM_EXECUTOR_INTERNAL_H_
#define TENSORFLOW_C_EXPERIMENTAL_STREAM_EXECUTOR_STREAM_EXECUTOR_INTERNAL_H_

#include "tensorflow/c/experimental/stream_executor/stream_executor.h"

    Op<TFRT_GPU_Dialect, mnemonic, traits> {
}

// TODO(b/260267885): We may add a device argument when we want to support
// GPU MIG.
def TransferToDeviceOp: Gpu_Op<"transfer_to_device"> {
  let summary = "Transfer a CPU tensor to device.";

  let description = [{
    Transfer a CPU tensor to device.

    Example:
      %device_tensor = gpurt.transfer_to_device %cpu_tensor
  }];