#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"

#include "llvm/Support/ErrorHandling.h"

#define DEBUG_TYPE "PimCompilerOptions"

namespace onnx_mlir {

llvm::cl::opt<PimEmissionTargetType> pimEmissionTarget(
  llvm::cl::desc("[Optional] Choose PIM-related target to emit (once selected it will cancel the other targets):"),
  llvm::cl::values(clEnumVal(EmitSpatial, "Lower model to spatial IR")),
  llvm::cl::values(clEnumVal(EmitPim, "Lower model to PIM IR")),
  llvm::cl::values(clEnumVal(EmitPimBufferized, "Lower model to PIM IR and bufferize it")),
  llvm::cl::values(clEnumVal(EmitPimCodegen, "Lower model to PIM IR and generate code for PIM")),
  llvm::cl::init(EmitPimCodegen),
  llvm::cl::cat(OnnxMlirOptions));

llvm::cl::opt<PimMergeSchedulerType> pimMergeScheduler(
  "pim-merge-scheduler",
  llvm::cl::desc("Scheduler used by the Spatial merge-compute-nodes pass"),
  llvm::cl::values(clEnumValN(MergeSchedulerPeft, "peft", "Use PEFT scheduling")),
  llvm::cl::values(clEnumValN(MergeSchedulerDcp, "dcp", "Use the legacy DCP-inspired scheduler")),
  llvm::cl::init(MergeSchedulerPeft),
  llvm::cl::cat(OnnxMlirOptions));

llvm::cl::opt<bool>
  pimOnlyCodegen("pim-only-codegen",
                 llvm::cl::desc("Only generate code for PIM (assume input is already in bufferized PIM IR)"),
                 llvm::cl::init(false),
                 llvm::cl::cat(OnnxMlirOptions));

llvm::cl::opt<bool> useExperimentalConvImpl("use-experimental-conv-impl",
                                            llvm::cl::desc("Use experimental implementation for convolution"),
                                            llvm::cl::init(false),
                                            llvm::cl::cat(OnnxMlirOptions));

llvm::cl::opt<bool> pimEmitJson("pim-emit-json",
                                llvm::cl::desc("Also emit per-core JSON instruction files alongside binary .pim files"),
                                llvm::cl::init(false),
                                llvm::cl::cat(OnnxMlirOptions));

llvm::cl::opt<size_t>
  crossbarSize("crossbar-size", llvm::cl::desc("Width and height of a single crossbar"), llvm::cl::init(2));

llvm::cl::opt<size_t>
  crossbarCountInCore("crossbar-count", llvm::cl::desc("Number of crossbars in each core"), llvm::cl::init(256));

llvm::cl::opt<long> coresCount("core-count",
                               llvm::cl::desc("Number of cores in the chip. Required for PIM compilation."),
                               llvm::cl::init(-1));

llvm::cl::opt<size_t> dcpCriticalWindowSize(
  "dcp-critical-window-size",
  llvm::cl::desc("Number of lowest-slack virtual nodes considered by each DCP coarsening iteration. "
                 "Use 0 to run the legacy full-graph DCP analysis. Only used by the DCP scheduler."),
  llvm::cl::init(4000));

llvm::cl::opt<bool>
  ignoreConcatError("ignore-concat-error",
                    llvm::cl::desc("Ignore ConcatOp corner case: do not assert and do a simplification"),
                    llvm::cl::init(false));

bool hasExplicitPimCoreCount() { return coresCount.getNumOccurrences() != 0; }

void verifyExplicitPimCoreCount() {
  if (!hasExplicitPimCoreCount())
    llvm::report_fatal_error("PIM compilation requires an explicit --core-count=<positive integer>");
  if (coresCount.getValue() <= 0)
    llvm::report_fatal_error("PIM compilation requires --core-count to be a positive integer");
}

} // namespace onnx_mlir