#include <cstdio>

#include <cstdlib>

#include <iostream>

#include <memory>

#include <string>

#include "mlir/Dialect/Arith/IR/Arith.h"

#include "mlir/Dialect/Linalg/IR/Linalg.h"

#include "mlir/Dialect/LLVMIR/LLVMDialect.h"

#include "mlir/Dialect/Func/IR/FuncOps.h"

#include "mlir/Dialect/Math/IR/Math.h"

#include "mlir/Dialect/Tensor/IR/Tensor.h"

#include "dialect/MLIRGen.h"

#include "dialect/NNDialect.h"

#include "mlir/IR/OwningOpRef.h"

#include "parser/parser.h"

#include "passes/Passes.h"

#include "mlir/Dialect/Bufferization/Transforms/Passes.h"

#include "mlir/Dialect/MemRef/Transforms/Passes.h"

#include "mlir/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.h"

#include "mlir/IR/BuiltinOps.h"

#include "mlir/IR/MLIRContext.h"

#include "mlir/Pass/PassManager.h"

#include "llvm/IR/LLVMContext.h"

#include "llvm/Support/raw_ostream.h"

#include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"

void printDebugInfo(std::string text, mlir::ModuleOp module) {

#ifdef DEBUG_MODE

llvm::outs() << text;

module->print(llvm::outs());

llvm::outs() << "\n";

#endif

}

int main(int argc, char **argv) {

std::string target_device = "cpu";

if (argc == 2) {

pFile = fopen(argv[1], "r");

if (pFile == NULL)

perror("Error opening file");

} else if (argc == 3) {

pFile = fopen(argv[1], "r");

if (pFile == NULL)

perror("Error opening file");

target_device = std::string(argv[2]);

} else {

std::cout << "Usage: ./mini_nn InputFile\n";

return 1;

}

llvm::outs() << "Target device: " << target_device << "\n";

// Extract directory path from input file for output files

std::string input_file_path(argv[1]);

std::string location_prefix = "";

size_t last_slash = input_file_path.find_last_of('/');

if (last_slash != std::string::npos) {

location_prefix =

input_file_path.substr(0, last_slash + 1); // Include the trailing slash

}

// Extract base filename without extension for naming output files

std::string base_filename = input_file_path;

if (last_slash != std::string::npos) {

base_filename = input_file_path.substr(last_slash + 1);

}

size_t dot_pos = base_filename.find_last_of('.');

if (dot_pos != std::string::npos) {

base_filename = base_filename.substr(0, dot_pos);

}

#ifdef DEBUG_MODE

std::cout << "Input file: " << input_file_path << std::endl;

std::cout << "Output directory: "

<< (location_prefix.empty() ? "./" : location_prefix) << std::endl;

std::cout << "Base filename: " << base_filename << std::endl;

#endif

// Parse the input program

Parser parser;

auto program = parser.parseProgram();

fclose(pFile);

if (!program) {

std::cerr << "Failed to parse program" << std::endl;

return 1;

}

#ifdef DEBUG_MODE

std::cout << "=== AST ===\n";

std::cout << program->toString() << std::endl;

#endif

// Create MLIR context and register all necessary dialects

mlir::MLIRContext context;

// Load specific dialects we need, without loadAllAvailableDialects() which

// might cause issues

context.getOrLoadDialect<mlir::nn::NNDialect>();

context.getOrLoadDialect<mlir::arith::ArithDialect>();

context.getOrLoadDialect<mlir::linalg::LinalgDialect>();

context.getOrLoadDialect<mlir::tensor::TensorDialect>();

context.getOrLoadDialect<mlir::math::MathDialect>();

context.getOrLoadDialect<mlir::func::FuncDialect>();

context.getOrLoadDialect<mlir::LLVM::LLVMDialect>();

// Generate MLIR from AST

auto module = mlir::nn::mlirGen(context, *program);

if (!module) {

std::cerr << "Failed to generate MLIR" << std::endl;

return 1;

}

printDebugInfo("\n=== MLIR (Before Shape Inference) ===\n", *module);

// Create a pass manager and add the shape inference pass

mlir::PassManager pm(&context);

// Since ShapeInferencePass operates on FuncOp, we need to use a nested pass

// manager

pm.addNestedPass<mlir::nn::FuncOp>(mlir::nn::createShapeInferencePass());

// Run the shape inference pass

if (mlir::failed(pm.run(*module))) {

std::cerr << "Failed to run shape inference pass" << std::endl;

return 1;

}

printDebugInfo("\n=== MLIR (After Shape Inference) ===\n", *module);

// Add lowering to Linalg pass

mlir::PassManager loweringPM(&context);

loweringPM.addPass(mlir::nn::createLowerToLinalgPass());

// Run the lowering pass

if (mlir::failed(loweringPM.run(*module))) {

std::cerr << "Failed to run lowering pass" << std::endl;

return 1;

}

printDebugInfo("\n=== MLIR (After Lowering to Linalg) ===\n", *module);

// Save the Linalg IR to a file for external processing

std::string linalg_file = location_prefix + base_filename + "_linalg.mlir";

std::error_code EC;

llvm::raw_fd_ostream linalg_output(linalg_file, EC);

if (EC) {

std::cerr << "Failed to open file for writing: " << linalg_file

<< std::endl;

return 1;

}

module->print(linalg_output);

linalg_output.close();

std::cout << "Saved Linalg IR to: " << linalg_file << std::endl;

mlir::PassManager affinePM(&context);

affinePM.addPass(mlir::nn::createLowerToLoopsPass());

// Run the LLVM lowering pass

if (mlir::failed(affinePM.run(*module))) {

std::cerr << "Failed to run LLVM lowering pass" << std::endl;

return 1;

}

std::string affine_file = location_prefix + base_filename + "_affine.mlir";

llvm::raw_fd_ostream affine_output(affine_file, EC);

if (EC) {

std::cerr << "Failed to open file for writing: " << affine_file

<< std::endl;

return 1;

}

module->print(affine_output);

affine_output.close();

std::cout << "Saved Affine IR to: " << affine_file << std::endl;

printDebugInfo("\n=== MLIR (After Lowering to LLVM) ===\n", *module);

mlir::PassManager llvmPM(&context);

if (target_device == "cpu") {

llvmPM.addPass(mlir::nn::createLowerLoopsToCPUPass());

} else {

llvmPM.addPass(mlir::nn::createLowerLoopsToGPUPass());

}

// Run the LLVM lowering pass

if (mlir::failed(llvmPM.run(*module))) {

std::cerr << "Failed to run LLVM lowering pass" << std::endl;

return 1;

}

std::string llvm_file = location_prefix + base_filename + "_llvm.mlir";

llvm::raw_fd_ostream llvm_output(llvm_file, EC);

if (EC) {

std::cerr << "Failed to open file for writing: " << llvm_file

<< std::endl;

return 1;

}

module->print(llvm_output);

llvm_output.close();

std::cout << "Saved Affine IR to: " << llvm_file << std::endl;

printDebugInfo("\n=== MLIR (After Lowering to LLVM) ===\n", *module);

return 0;

}