//go:build s390x && linux

package unix

import (
	"unsafe"
)

//sys	EpollWait(epfd int, events []EpollEvent, msec int) (n int, err error)
//sys	Fadvise(fd int, offset int64, length int64, advice int) (err error) = SYS_FADVISE64
//sys	Fchown(fd int, uid int, gid int) (err error)
//sys	Fstat(fd int, stat *Stat_t) (err error)
//sys	Fstatat(dirfd int, path string, stat *Stat_t, flags int) (err error) = SYS_NEWFSTATAT

// 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>
}

      auto device = op->getAttrOfType<StringAttr>(kDevice);
      for (auto *operand : operands)
        propagateIfChanged(operand, operand->SetDevice(device));
    } else {
      // Propagate device through other ops. These ops might have their
      // own device annotation, but that's fine. We only care about
      // where the TPUExecute ops live.
      StringAttr device;
      for (const Device *d : results) {

      *this, "device-specs",
      llvm::cl::desc(
          "comma separated list of device specs, like CPU, GPU, Hexagon."),
      llvm::cl::ZeroOrMore};

  void getDependentDialects(mlir::DialectRegistry& registry) const override {
    if (!module_) {
      for (const auto& device : device_specs_flag_) {
        auto* hardware = this->GetTargetHardware(device);
        if (hardware == nullptr) continue;

  return name;
}

// Get the target annotation form the op.
inline std::optional<std::string> GetTargetAnnotation(Operation* op) {
  auto device = op->getAttrOfType<StringAttr>(kDevice);
  if (device == nullptr || device.getValue().empty()) return std::nullopt;

  return GetCanonicalHardwareName(device.getValue().str());
}

// Get inference type attribute from the operation if available.

// Creates a pass that extracts outside compilation (Host ops inside device
// cluster) at head/tail of Device cluster to run before/after XLA computation.
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateExtractHeadTailOutsideCompilationPass();

// Creates a pass that extract outside compilation (Host ops inside cevice
// cluster) ops to a separate parallel_execute region to run on CPU.

    return std::nullopt;

  if (!HasValidHardwareTarget(op)) return std::nullopt;

  auto device = op->getAttrOfType<mlir::StringAttr>(mlir::TFL::tac::kDevice);
  if (device == nullptr) return std::nullopt;

  llvm::StringRef device_name_str = device.getValue();
  return device_name_str.str();
}

std::optional<std::vector<float>> GetPerDeviceCosts(

                          XlaCompileOnDemandOp);                               \
  REGISTER_KERNEL_BUILDER(Name("XlaSvd").Device(DEVICE),                       \
                          XlaCompileOnDemandOp);                               \
  REGISTER_KERNEL_BUILDER(Name("XlaDot").Device(DEVICE),                       \
                          XlaCompileOnDemandOp);                               \

                                       mlir::Builder* builder) {
  // Parse GPU device compute capability from physical device description.
  static auto* r = new llvm::Regex("compute capability: ([0-9]+)\\.([0-9]+)");

  llvm::SmallVector<llvm::StringRef, 3> cc;
  if (r->match(device.attributes().physical_device_desc(), &cc)) {
    return mlir::TF::GpuDeviceMetadata::get(