Add option to compiler to Preprocessing step. (#12140)
Add an option to `--compile-to` which allows stopping the compilation after preprocessing.
Usage : `iree-compile --compile-to=preprocessing`.
diff --git a/compiler/src/iree/compiler/Pipelines/Pipelines.cpp b/compiler/src/iree/compiler/Pipelines/Pipelines.cpp
index 180ef39..904f297 100644
--- a/compiler/src/iree/compiler/Pipelines/Pipelines.cpp
+++ b/compiler/src/iree/compiler/Pipelines/Pipelines.cpp
@@ -126,7 +126,11 @@
// No flow/stream processing (implies no tensors).
break;
default:
+ IREE_TRACE_ADD_BEGIN_FRAME_PASS(passManager, "Preprocessing");
IREE::buildPreprocessingPassPipeline(passManager, preprocessingOptions);
+ IREE_TRACE_ADD_END_FRAME_PASS(passManager, "Preprocessing");
+ if (compileTo == IREEVMPipelinePhase::Preprocessing)
+ return; // early-exit
IREE_TRACE_ADD_BEGIN_FRAME_PASS(passManager, "Flow");
IREE::Flow::buildFlowTransformPassPipeline(passManager, flowOptions);
diff --git a/compiler/src/iree/compiler/Pipelines/Pipelines.h b/compiler/src/iree/compiler/Pipelines/Pipelines.h
index 319fe54..c8781b2 100644
--- a/compiler/src/iree/compiler/Pipelines/Pipelines.h
+++ b/compiler/src/iree/compiler/Pipelines/Pipelines.h
@@ -31,6 +31,7 @@
enum class IREEVMPipelinePhase {
Input,
ABI,
+ Preprocessing,
Flow,
Stream,
HAL,
@@ -47,6 +48,8 @@
"input dialects (linalg/etc).");
callback(IREEVMPipelinePhase::ABI, "abi",
"Adjusts program ABI for the specified execution environment.");
+ callback(IREEVMPipelinePhase::Preprocessing, "preprocessing",
+ "Compiles up to the `preprocessing` specified");
callback(IREEVMPipelinePhase::Flow, "flow",
"Compiles up to the `flow` dialect.");
callback(IREEVMPipelinePhase::Stream, "stream",
diff --git a/tests/compiler_driver/preprocessing_flags.mlir b/tests/compiler_driver/preprocessing_flags.mlir
index e2cb536..ecd41db 100644
--- a/tests/compiler_driver/preprocessing_flags.mlir
+++ b/tests/compiler_driver/preprocessing_flags.mlir
@@ -1,4 +1,4 @@
-// RUN: iree-compile --iree-hal-target-backends=llvm-cpu --output-format=vm-asm \
+// RUN: iree-compile --iree-hal-target-backends=llvm-cpu --compile-to=preprocessing \
// RUN: --iree-preprocessing-pass-pipeline="builtin.module(func.func(iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=16}))" \
// RUN: --mlir-print-ir-after=iree-preprocessing-convert-conv2d-to-img2col --mlir-print-ir-after=iree-preprocessing-pad-linalg-ops %s 2>&1 \
// RUN: | FileCheck %s
@@ -9,17 +9,19 @@
return %0 : tensor<10x30xf32>
}
// Just check that the pass runs, and that the compilation finishes
-// CHECK: ConvertConv2DToImg2Col (iree-preprocessing-convert-conv2d-to-img2col)
-// CHECK: PadLinalgOps (iree-preprocessing-pad-linalg-ops)
-// CHECK-DAG: %[[ARG0:.+]] = hal.tensor.import %{{[a-zA-Z0-9]+}} : !hal.buffer_view -> tensor<10x20xf32>
-// CHECK-DAG: %[[ARG1:.+]] = hal.tensor.import %{{[a-zA-Z0-9]+}} : !hal.buffer_view -> tensor<20x30xf32>
-// CHECK-DAG: %[[ARG2:.+]] = hal.tensor.import %{{[a-zA-Z0-9]+}} : !hal.buffer_view -> tensor<10x30xf32>
-// CHECK-DAG: %[[PAD0:.+]] = tensor.pad %[[ARG0]] low[0, 0] high[6, 12]
-// CHECK-DAG: %[[PAD1:.+]] = tensor.pad %[[ARG1]] low[0, 0] high[12, 2]
-// CHECK-DAG: %[[PAD2:.+]] = tensor.pad %[[ARG2]] low[0, 0] high[6, 2]
-// CHECK: %[[PADDED:.+]] = linalg.matmul
-// CHECK-SAME: ins(%[[PAD0]], %[[PAD1]] :
-// CHECK-SAME: outs(%[[PAD2]] :
-// CHECK: %[[SLICE:.+]] = tensor.extract_slice %[[PADDED]][0, 0] [10, 30] [1, 1]
-// CHECK: hal.tensor.export %[[SLICE]]
-// CHECK: vm.module
+// CHECK: ConvertConv2DToImg2Col (iree-preprocessing-convert-conv2d-to-img2col)
+// CHECK: PadLinalgOps (iree-preprocessing-pad-linalg-ops)
+// CHECK-LABEL: module
+// CHECK-NEXT: func.func @test(
+// CHECK-DAG: %[[ARG0:.+]] = hal.tensor.import %{{[a-zA-Z0-9]+}} : !hal.buffer_view -> tensor<10x20xf32>
+// CHECK-DAG: %[[ARG1:.+]] = hal.tensor.import %{{[a-zA-Z0-9]+}} : !hal.buffer_view -> tensor<20x30xf32>
+// CHECK-DAG: %[[ARG2:.+]] = hal.tensor.import %{{[a-zA-Z0-9]+}} : !hal.buffer_view -> tensor<10x30xf32>
+// CHECK-DAG: %[[PAD0:.+]] = tensor.pad %[[ARG0]] low[0, 0] high[6, 12]
+// CHECK-DAG: %[[PAD1:.+]] = tensor.pad %[[ARG1]] low[0, 0] high[12, 2]
+// CHECK-DAG: %[[PAD2:.+]] = tensor.pad %[[ARG2]] low[0, 0] high[6, 2]
+// CHECK: %[[PADDED:.+]] = linalg.matmul
+// CHECK-SAME: ins(%[[PAD0]], %[[PAD1]] :
+// CHECK-SAME: outs(%[[PAD2]] :
+// CHECK: %[[SLICE:.+]] = tensor.extract_slice %[[PADDED]][0, 0] [10, 30] [1, 1]
+// CHECK: hal.tensor.export %[[SLICE]]
+// CHECK-NOT: vm.module
