#define TENSORFLOW_COMPILER_MLIR_REGISTER_COMMON_DIALECTS_H_

#include "mlir/IR/DialectRegistry.h"  // from @llvm-project

namespace mlir {

// Inserts common Tensorflow dialects used for offline tools.
void RegisterCommonToolingDialects(mlir::DialectRegistry& registry);

};  // namespace mlir

#include <gtest/gtest.h>
#include "mlir/IR/DialectRegistry.h"  // from @llvm-project

namespace mlir {
namespace {

TEST(RegisterCommonDialectsTest, DoesntCrash) {
  mlir::DialectRegistry registry;

  mlir::RegisterCommonToolingDialects(registry);

  EXPECT_FALSE(registry.getDialectNames().empty());
}

}  // namespace

using ::mlir::OwningOpRef;
using ::tsl::StatusOr;

absl::StatusOr<OwningOpRef<ModuleOp>> GetMlirModuleFromString(
    absl::string_view module_string, MLIRContext* context) {
  DialectRegistry mlir_registry;
  RegisterCommonToolingDialects(mlir_registry);
  context->appendDialectRegistry(mlir_registry);

  OwningOpRef<ModuleOp> mlir_module;
  auto status =
      tensorflow::DeserializeMlirModule(module_string, context, &mlir_module);

  return tensorflow::GetDataDependencyFilepath(
      "tensorflow/compiler/mlir/tf2xla/internal/testdata/");
}

class LoggingHooksTest : public ::testing::Test {
 public:
  LoggingHooksTest() {
    mlir::RegisterCommonToolingDialects(registry_);
    context_.appendDialectRegistry(registry_);
    context_.loadAllAvailableDialects();
    env_ = Env::Default();

    test_group_name_ = "TestGroup";
    test_dir_ = testing::TmpDir();

    populateWithGenerated(patterns);
  }
};

}  // namespace

int main(int argc, char *argv[]) {
  tensorflow::InitMlir y(&argc, &argv);

  mlir::DialectRegistry registry;
  mlir::RegisterCommonToolingDialects(registry);

  registry.addExtension(
      +[](mlir::MLIRContext *ctx, mlir::TF::TensorFlowDialect *dialect) {
        dialect->addInterfaces<TFReductionPatternInterface>();
      });

    count++;
    idx = str.find(substr, idx + 1);
  }
  return count;
}

class TensorflowDialectToExecutorTest : public ::testing::Test {
 public:
  TensorflowDialectToExecutorTest() {
    mlir::RegisterCommonToolingDialects(registry_);
    context_.appendDialectRegistry(registry_);
    context_.loadAllAvailableDialects();
  }

  absl::Status CreateMlirModule(std::string mlir_module_filename) {

#include "tensorflow/compiler/mlir/tools/kernel_gen/ir/tf_framework_ops.h"
#include "xla/mlir/framework/ir/xla_framework.h"
#include "xla/mlir_hlo/mhlo/IR/register.h"

namespace mlir {

void RegisterCommonToolingDialects(mlir::DialectRegistry& registry) {
  mlir::RegisterAllTensorFlowDialects(registry);
  mlir::mhlo::registerAllMhloDialects(registry);
  mlir::registerAllDialects(registry);
  mlir::registerAllExtensions(registry);

// correctness test versus a test to easily ensure the op registry is kept
// in sync.
TEST_F(LegalizationOpConfigTest, CountAllMlirLoweringPatterns) {
  DialectRegistry dialect_registry;
  mlir::RegisterCommonToolingDialects(dialect_registry);

  MLIRContext context(dialect_registry);
  context.loadAllAvailableDialects();

  RewritePatternSet mlir_legalize_lower_patterns(&context);

      "tensorflow/compiler/mlir/tf2xla/api/v1/testdata/");
}

class SessionClusterTensorflowDialectTest : public ::testing::Test {
 public:
  SessionClusterTensorflowDialectTest() {
    mlir::RegisterCommonToolingDialects(registry_);
    context_.appendDialectRegistry(registry_);
    context_.loadAllAvailableDialects();
  }

  absl::Status CreateMlirModule(std::string mlir_module_filename) {

      "tensorflow/compiler/mlir/tf2xla/api/v1/testdata/");
}

class TensorflowDialectToExecutorTest : public ::testing::Test {
 public:
  TensorflowDialectToExecutorTest() {
    mlir::RegisterCommonToolingDialects(registry_);
    context_.appendDialectRegistry(registry_);
    context_.loadAllAvailableDialects();
  }

  absl::Status CreateMlirModule(std::string mlir_module_filename) {