graph_compiler.cpp

#include "mlir/IR/AsmState.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/MLIRContext.h"
#include "mlir/Parser/Parser.h"
#include "mlir/Pass/PassManager.h"
#include "mlir/Support/FileUtilities.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/SourceMgr.h"

#include "Gateway.h"

namespace gateway {
void registerGraphToKernelPass();
}

int main(int argc, char **argv) {
    mlir::registerAsmPrinterCLOptions();
    mlir::registerMLIRContextCLOptions();
    mlir::registerPassManagerCLOptions();

    llvm::cl::opt<std::string> inputFilename(llvm::cl::Positional, llvm::cl::desc("<input file>"), llvm::cl::init("-"));

    llvm::cl::ParseCommandLineOptions(argc, argv, "Graph Compiler\n");

    mlir::MLIRContext context;
    context.getOrLoadDialect<gateway::GatewayDialect>();
    context.getOrLoadDialect<mlir::func::FuncDialect>();
    context.getOrLoadDialect<mlir::arith::ArithDialect>();
    context.getOrLoadDialect<mlir::memref::MemRefDialect>();

    mlir::OwningOpRef<mlir::ModuleOp> module;
    llvm::SourceMgr sourceMgr;
    mlir::SourceMgrDiagnosticHandler sourceMgrHandler(sourceMgr, &context);

    std::string errorMessage;
    auto file = mlir::openInputFile(inputFilename, &errorMessage);
    if (!file) {
        llvm::errs() << errorMessage << "\n";
        return 1;
    }

    sourceMgr.AddNewSourceBuffer(std::move(file), llvm::SMLoc());
    module = mlir::parseSourceFile<mlir::ModuleOp>(sourceMgr, &context);
    if (!module) {
        llvm::errs() << "Error can't load file " << inputFilename << "\n";
        return 1;
    }

    mlir::PassManager pm(&context);
    gateway::registerGraphToKernelPass();
    pm.addPass(mlir::createModuleToFunctionPassAdaptor(gateway::createGraphToKernelPass()));

    if (mlir::failed(pm.run(*module))) {
        llvm::errs() << "Error running pass manager\n";
        return 1;
    }

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

    return 0;
}