heir/lib/Transforms/AnnotateModule/AnnotateModule.cpp at cc04f333881e3293b78ead8066afc4d430c79e4a · google/heir · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
#include "lib/Transforms/AnnotateModule/AnnotateModule.h"

#include "lib/Dialect/ModuleAttributes.h"
#include "lib/Target/CompilationTarget/CompilationTarget.h"
#include "mlir/include/mlir/IR/BuiltinOps.h"  // from @llvm-project
#include "mlir/include/mlir/Support/LLVM.h"   // from @llvm-project

namespace mlir {
namespace heir {

#define GEN_PASS_DEF_ANNOTATEMODULE
#include "lib/Transforms/AnnotateModule/AnnotateModule.h.inc"

struct AnnotateModule : impl::AnnotateModuleBase<AnnotateModule> {
  using AnnotateModuleBase::AnnotateModuleBase;

  void runOnOperation() override {
    ModuleOp module = cast<ModuleOp>(getOperation());

    if (scheme == "bgv") {
      moduleSetBGV(module);
    } else if (scheme == "bfv") {
      moduleSetBFV(module);
    } else if (scheme == "ckks") {
      moduleSetCKKS(module);
    } else if (scheme == "cggi") {
      moduleSetCGGI(module);
    }

    if (!backend.empty()) {
      if (!CompilationTargetRegistry::get(backend)) {
        module.emitError() << "Unknown backend: " << backend;
        signalPassFailure();
        return;
      }

      if (backend == "openfhe") {
        moduleSetOpenfhe(module);
      } else if (backend == "lattigo") {
        moduleSetLattigo(module);
      }
    }
  }
};

}  // namespace heir
}  // namespace mlir