}
  filename = llvm::Twine(filename).concat(".mlir").str();
  return filename;
}

// Simple raw_ostream that prints to stderr.
struct LogInfoRawStream : public llvm::raw_ostream {
  LogInfoRawStream() { SetUnbuffered(); }
  ~LogInfoRawStream() override = default;
  uint64_t current_pos() const override { return 0; }

  void write_impl(const char* ptr, size_t size) override {
    fprintf(stderr, "%.*s", static_cast<int>(size), ptr);

namespace {

// Simple raw_ostream that prints to a file (doesn't take ownership).
struct WritableFileRawStream : public llvm::raw_ostream {
  explicit WritableFileRawStream(WritableFile* file) : file(file) {
    SetUnbuffered();
  }
  ~WritableFileRawStream() override = default;
  uint64_t current_pos() const override { return 0; }

  void write_impl(const char* ptr, size_t size) override {

namespace {
static const char* kTestInstrumentationName = "test-intrumentatron";
static const char* kTestInstrumentationSearch = "tf.Identity";

struct StringStream : public llvm::raw_ostream {
  StringStream() { SetUnbuffered(); }
  ~StringStream() override = default;
  uint64_t current_pos() const override { return 0; }

  void write_impl(const char* ptr, size_t size) override {
    ss.write(ptr, size);
  }
  std::stringstream ss;

namespace {

// Simple raw_ostream that prints to a file.
struct WritableFileRawStream : public llvm::raw_ostream {
  explicit WritableFileRawStream(std::unique_ptr<tsl::WritableFile> file)
      : file(std::move(file)) {
    SetUnbuffered();
  }
  ~WritableFileRawStream() override = default;

  uint64_t current_pos() const override {
    int64_t position;
    if (file->Tell(&position).ok()) {
      return position;
    } else {