//===---------------------------------------------------------------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===---------------------------------------------------------------------===//
//
// This file implements a pass to raise operations to arith dialect.
//===---------------------------------------------------------------------===//

#ifndef DISABLE_POLYMER
#include "../polymer/mlir/include/mlir/Conversion/Polymer/Support/IslScop.h"
#include "../polymer/mlir/include/mlir/Conversion/Polymer/Target/ISL.h"
#endif
#include "Enzyme/MLIR/Dialect/Dialect.h"
#include "Enzyme/MLIR/Dialect/Ops.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Complex/IR/Complex.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Math/IR/Math.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/IRMapping.h"
#include "src/enzyme_ad/jax/Dialect/Dialect.h"
#include "src/enzyme_ad/jax/Dialect/Ops.h"
#include "src/enzyme_ad/jax/Passes/Passes.h"
#include "llvm/Support/DebugLog.h"

#define DEBUG_TYPE "remove-atomics"

namespace mlir {
namespace enzyme {
#define GEN_PASS_DEF_TESTPOLYMERPASS
#include "src/enzyme_ad/jax/Passes/Passes.h.inc"
} // namespace enzyme
} // namespace mlir

using namespace mlir;
using namespace mlir::enzyme;

namespace {
struct TestPolymerPass
    : public enzyme::impl::TestPolymerPassBase<TestPolymerPass> {
  using TestPolymerPassBase::TestPolymerPassBase;
  void runOnOperation() override {
#ifndef DISABLE_POLYMER
    getOperation()->walk([](Operation *op) {
      if (!isa<func::FuncOp, enzymexla::GPUWrapperOp>(op))
        return;
      LDBG() << "Processing " << op;
      llvm::errs() << "Processing " << op->getName() << "\n";
      std::unique_ptr<polymer::IslScop> scop = polymer::createIslFromFuncOp(op);
      if (!scop) {
        llvm::errs() << "Failed to build scop\n";
        return;
      }
      if (scop->buildSchedule().succeeded()) {
        llvm::errs() << "Schedule:\n";
        isl_schedule_dump(scop->getScheduleTree().get());
        llvm::errs() << "Accesses:\n";
        scop->dumpAccesses(llvm::errs());
      } else {
        llvm::errs() << "Failed to build schedule\n";
      }
    });
#endif
  }
};
} // namespace
