Fixing the --compile-to= flag and adding a test of all phases. (#11345)
Fixes #11330.
diff --git a/compiler/src/iree/compiler/Pipelines/Pipelines.h b/compiler/src/iree/compiler/Pipelines/Pipelines.h
index d3d5951..f149978 100644
--- a/compiler/src/iree/compiler/Pipelines/Pipelines.h
+++ b/compiler/src/iree/compiler/Pipelines/Pipelines.h
@@ -42,8 +42,6 @@
inline static void enumerateIREEVMPipelinePhases(
std::function<void(IREEVMPipelinePhase, StringRef name, StringRef desc)>
callback) {
- callback(IREEVMPipelinePhase::End, "end",
- "Complete the full compilation pipeline.");
callback(IREEVMPipelinePhase::Input, "input",
"Performs input processing and lowering into core IREE "
"input dialects (linalg/etc).");
@@ -56,6 +54,8 @@
callback(IREEVMPipelinePhase::HAL, "hal",
"Compiles up to the `hal` dialect, including codegen.");
callback(IREEVMPipelinePhase::VM, "vm", "Compiles up to the `vm` dialect.");
+ callback(IREEVMPipelinePhase::End, "end",
+ "Complete the full compilation pipeline.");
}
// Builds a pass pipeline to perform end-to-end compilation from a
diff --git a/compiler/src/iree/compiler/Tools/iree_compile_lib.cc b/compiler/src/iree/compiler/Tools/iree_compile_lib.cc
index fe7dc93..37884b3 100644
--- a/compiler/src/iree/compiler/Tools/iree_compile_lib.cc
+++ b/compiler/src/iree/compiler/Tools/iree_compile_lib.cc
@@ -93,14 +93,16 @@
llvm::cl::desc("Verifies the IR for correctness throughout compilation."),
llvm::cl::init(true));
- llvm::cl::opt<llvm::StringRef> compileTo(
+ llvm::cl::opt<IREEVMPipelinePhase> compileTo(
"compile-to",
llvm::cl::desc(
"Compilation phase to run up until before emitting output."),
- llvm::cl::init("end"));
+ llvm::cl::init(IREEVMPipelinePhase::End));
+ SmallVector<std::string> compileToPhases;
enumerateIREEVMPipelinePhases(
[&](IREEVMPipelinePhase phase, StringRef name, StringRef desc) {
- compileTo.getParser().addLiteralOption(name, name, desc);
+ compileTo.getParser().addLiteralOption(name, phase, desc);
+ compileToPhases.push_back(name.str());
});
// Misc options.
@@ -178,8 +180,8 @@
InvState r(s);
ireeCompilerInvocationEnableConsoleDiagnostics(r.inv);
- ireeCompilerInvocationSetCompileToPhase(r.inv,
- std::string(compileTo).c_str());
+ ireeCompilerInvocationSetCompileToPhase(
+ r.inv, compileToPhases[static_cast<int>(compileTo.getValue())].c_str());
ireeCompilerInvocationSetVerifyIR(r.inv, verifyIR);
if (!ireeCompilerInvocationParseSource(r.inv, source)) return false;
@@ -203,7 +205,7 @@
}
// Ending early and just emitting IR.
- if (compileTo != "end") {
+ if (compileTo != IREEVMPipelinePhase::End) {
if (auto error = ireeCompilerInvocationOutputIR(r.inv, s.output)) {
s.handleError(error);
return false;
diff --git a/tools/test/BUILD b/tools/test/BUILD
index 8ea1678..827766b 100644
--- a/tools/test/BUILD
+++ b/tools/test/BUILD
@@ -18,6 +18,7 @@
name = "lit",
srcs = enforce_glob(
[
+ "compile_to_phase.mlir",
"executable_benchmarks.mlir",
"iree-benchmark-module.mlir",
"iree-run-mlir.mlir",
diff --git a/tools/test/CMakeLists.txt b/tools/test/CMakeLists.txt
index cb854ac..dcb3087 100644
--- a/tools/test/CMakeLists.txt
+++ b/tools/test/CMakeLists.txt
@@ -14,6 +14,7 @@
NAME
lit
SRCS
+ "compile_to_phase.mlir"
"executable_benchmarks.mlir"
"iree-benchmark-module.mlir"
"iree-run-mlir.mlir"
diff --git a/tools/test/compile_to_phase.mlir b/tools/test/compile_to_phase.mlir
new file mode 100644
index 0000000..cb3a8ed
--- /dev/null
+++ b/tools/test/compile_to_phase.mlir
@@ -0,0 +1,35 @@
+// RUN: iree-compile --compile-to=input --iree-hal-target-backends=vmvx %s | FileCheck %s --check-prefix=INPUT-PHASE
+// INPUT-PHASE: func.func @abs(%[[ARG0:.+]]: tensor<f32>)
+// INPUT-PHASE: math.absf %[[ARG0]] : tensor<f32>
+
+// RUN: iree-compile --compile-to=abi --iree-hal-target-backends=vmvx %s | FileCheck %s --check-prefix=ABI-PHASE
+// ABI-PHASE: func.func @abs(%[[ARG0:.+]]: !hal.buffer_view)
+// ABI-PHASE: %[[INPUT:.+]] = hal.tensor.import %[[ARG0]] : !hal.buffer_view -> tensor<f32>
+// ABI-PHASE: math.absf %[[INPUT]] : tensor<f32>
+
+// RUN: iree-compile --compile-to=flow --iree-hal-target-backends=vmvx %s | FileCheck %s --check-prefix=FLOW-PHASE
+// FLOW-PHASE: flow.executable.export public @abs_dispatch_0
+// FLOW-PHASE: flow.dispatch @abs_dispatch_0
+
+// RUN: iree-compile --compile-to=stream --iree-hal-target-backends=vmvx %s | FileCheck %s --check-prefix=STREAM-PHASE
+// STREAM-PHASE: stream.executable.export public @abs_dispatch_0
+// STREAM-PHASE: stream.cmd.dispatch @abs_dispatch_0
+
+// RUN: iree-compile --compile-to=hal --iree-hal-target-backends=vmvx %s | FileCheck %s --check-prefix=HAL-PHASE
+// HAL-PHASE: hal.executable private @abs_dispatch_0
+// HAL-PHASE: hal.executable.binary
+// HAL-PHASE: hal.command_buffer.dispatch
+
+// RUN: iree-compile --compile-to=vm --iree-hal-target-backends=vmvx %s | FileCheck %s --check-prefix=VM-PHASE
+// VM-PHASE: vm.rodata private @abs_dispatch_0
+// VM-PHASE: vm.call @hal.command_buffer.dispatch
+
+// RUN: iree-compile --output-format=vm-asm --compile-to=end --iree-hal-target-backends=vmvx %s | FileCheck %s --check-prefix=END-PHASE
+// RUN: iree-compile --output-format=vm-asm --iree-hal-target-backends=vmvx %s | FileCheck %s --check-prefix=END-PHASE
+// END-PHASE: vm.rodata private @abs_dispatch_0
+// END-PHASE: vm.call @hal.command_buffer.dispatch
+
+func.func @abs(%input : tensor<f32>) -> (tensor<f32>) {
+ %result = math.absf %input : tensor<f32>
+ return %result : tensor<f32>
+}
