// `tf.devices` dictionary attribute with a full device name as a key, and
// 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);

}

// Macro used to parse initializer list for tensorshape
#define DIMARGS(...) \
  { __VA_ARGS__ }
// // Random parameters for testing
constexpr char longTagParam[] = "LONGTAG____________________________";
constexpr float largeValueParam = 2352352.2623433;

#define BM_ScalarSummaryDev(device, dims, name, tag, value)                 \
  void BM_ScalarSummary##name##device(::testing::benchmark::State& state) { \

limitations under the License.
==============================================================================*/

#ifndef TENSORFLOW_C_EAGER_CUSTOM_DEVICE_TESTUTIL_H_
#define TENSORFLOW_C_EAGER_CUSTOM_DEVICE_TESTUTIL_H_

// A simple logging device to test custom device registration.
#include <memory>

#include "tensorflow/c/c_api.h"
#include "tensorflow/c/eager/c_api.h"
#include "tensorflow/c/eager/c_api_experimental.h"

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

  void CopyCPUTensorToDevice(const Tensor* cpu_tensor, Device* device,
                             Tensor* device_tensor, StatusCallback done,
                             bool sync_dst_compute) const override;
  void CopyDeviceTensorToCPU(const Tensor* device_tensor,
                             absl::string_view tensor_name, Device* device,
                             Tensor* cpu_tensor, StatusCallback done) override;

                             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:

  // PjRtTensorBuffer. The PjRtTensorBuffer holds the pointer to the device
  // memory. It also owns the PjRtBuffer. If output_tensor does not use
  // PjRtTensorBuffer and the opaque device memory is the same, update the
  // output_tensor->buf_ so that the same device memory will not be double-free.
  // Otherwise a new Tensor will be created with the PjRtTensorBuffer.
  //
  // TODO(b/289001822): Create a unit test to cover this function.

  }

  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

// Creates a pass that merges device variable reads/updates into the surrounded
// TPUExecute node. This allows the execute node to perform in-place variable
// updates.
std::unique_ptr<OperationPass<ModuleOp>>
CreateTPUMergeVariablesWithExecutePass();

#define GEN_PASS_REGISTRATION
#define GEN_PASS_DECL_TPUMERGEVARIABLESWITHEXECUTEPASS
#define GEN_PASS_DECL_TPUREWRITEPASS
#define GEN_PASS_DECL_TPUVARIABLERUNTIMEREFORMATTINGPASS