// RUN: tf-opt %s -tf-layout-assignment -verify-diagnostics | FileCheck %s --dump-input=always

module attributes {
  tf.devices = {"/device:GPU:0" = #tf_type.gpu_device_metadata<cc_major = 6, cc_minor = 0>}
} {

// CHECK-LABEL: func @transposeConv2D_3x3_f16
func.func @transposeConv2D_3x3_f16(%input: tensor<1x28x28x64xf16>, %filter: tensor<3x3x64x64xf16>) -> tensor<1x26x26x64xf16> {
  // cuDNN prefers NCHW data format for spatial convolutions in f16 before

// RUN: tf-opt %s -tf-layout-assignment -verify-diagnostics | FileCheck %s --dump-input=always

module attributes {
  tf.devices = {"/device:GPU:0" = #tf_type.gpu_device_metadata<cc_major = 7, cc_minor = 0>}
} {

// CHECK-LABEL: func @transposeConv2D_3x3_f32
func.func @transposeConv2D_3x3_f32(%input: tensor<1x28x28x64xf32>, %filter: tensor<3x3x64x64xf32>) -> tensor<1x26x26x64xf32> {
  // cuDNN prefers NCHW data format for spatial convolutions.

  void AddGpuDevice(const ParsedName& device,
                    const GpuDeviceMetadata& metadata);

  llvm::ArrayRef<ParsedName> device_names() const { return device_names_; }
  size_t NumDevices() const { return device_names_.size(); }

  // Returns GPU device metadata if it is available, otherwise returns None.
  std::optional<GpuDeviceMetadata> GetGpuDeviceMetadata(
      const ParsedName& device) const;

 private:

}

void RuntimeDevices::AddGpuDevice(const ParsedName& device,
                                  const GpuDeviceMetadata& metadata) {
  device_names_.push_back(device);
  gpu_metadata_.insert({DeviceNameUtils::ParsedNameToString(device), metadata});
}

std::optional<GpuDeviceMetadata> RuntimeDevices::GetGpuDeviceMetadata(
    const ParsedName& device) const {

  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(
        builder->getContext(), std::stoi(cc[1].str()), std::stoi(cc[2].str()));
  }

  return builder->getUnitAttr();
}

// Get devices from an array of string attributes.

// device metadata as a value.
//
// Device names added in full parsed device form:
//   /job:<name>/replica:<replica>/task:<task>/device:<type>:<device_num>
//
// Supported device metadata types:
// (1) GpuDeviceMetadata: GPU device compute capability.
void AddDevicesToOp(mlir::Operation* op, const DeviceSet* device_set);

// Collects devices information from an op `tf.devices` attributes. Returns

  ASSERT_NE(device_meta_0, nullptr);

  // GPU device successfully parsed compute capability from description.
  auto device_meta_1 =
      mlir::dyn_cast<mlir::TF::GpuDeviceMetadata>(devices_attr.get(gpu0));
  ASSERT_NE(device_meta_1, nullptr);
  ASSERT_EQ(device_meta_1.getCcMajor(), 7);
  ASSERT_EQ(device_meta_1.getCcMinor(), 0);