feat(compiler): Add option --batch-concrete-ops and action dump-concrete-with-loops

The new option `--batch-concrete-ops` invokes the batching pass after
lowering to the Concrete dialect and after lowering linalg operations
with operations from the Concrete dialect to loops.

The new action `dump-concrete-with-loops` dumps the IR right before
batching.
This commit is contained in:
Andi Drebes
2022-11-10 16:55:16 +01:00
parent 75b70054b2
commit c9bb6541e9
6 changed files with 63 additions and 13 deletions

View File

@@ -48,6 +48,7 @@ enum Action {
DUMP_FHE,
DUMP_TFHE,
DUMP_CONCRETE,
DUMP_CONCRETEWITHLOOPS,
DUMP_BCONCRETE,
DUMP_STD,
DUMP_LLVM_DIALECT,
@@ -121,6 +122,9 @@ static llvm::cl::opt<enum Action> action(
"Lower to TFHE and dump result")),
llvm::cl::values(clEnumValN(Action::DUMP_CONCRETE, "dump-concrete",
"Lower to Concrete and dump result")),
llvm::cl::values(clEnumValN(
Action::DUMP_CONCRETEWITHLOOPS, "dump-concrete-with-loops",
"Lower to Concrete, replace linalg ops with loops and dump result")),
llvm::cl::values(
clEnumValN(Action::DUMP_BCONCRETE, "dump-bconcrete",
"Lower to Bufferized Concrete and dump result")),
@@ -162,6 +166,13 @@ llvm::cl::opt<bool> loopParallelize(
"Generate (and execute if JIT) parallel loops from Linalg operations"),
llvm::cl::init(false));
llvm::cl::opt<bool> batchConcreteOps(
"batch-concrete-ops",
llvm::cl::desc(
"Hoist scalar Concrete operations with corresponding batched "
"operations out of loop nests as batched operations"),
llvm::cl::init(false));
llvm::cl::opt<bool> dataflowParallelize(
"parallelize-dataflow",
llvm::cl::desc(
@@ -288,6 +299,7 @@ cmdlineCompilationOptions() {
options.autoParallelize = cmdline::autoParallelize;
options.loopParallelize = cmdline::loopParallelize;
options.dataflowParallelize = cmdline::dataflowParallelize;
options.batchConcreteOps = cmdline::batchConcreteOps;
options.optimizeConcrete = cmdline::optimizeConcrete;
options.emitGPUOps = cmdline::emitGPUOps;
@@ -457,6 +469,9 @@ mlir::LogicalResult processInputBuffer(
case Action::DUMP_CONCRETE:
target = mlir::concretelang::CompilerEngine::Target::CONCRETE;
break;
case Action::DUMP_CONCRETEWITHLOOPS:
target = mlir::concretelang::CompilerEngine::Target::CONCRETEWITHLOOPS;
break;
case Action::DUMP_BCONCRETE:
target = mlir::concretelang::CompilerEngine::Target::BCONCRETE;
break;