End-to-end example for lowering config driver lowering. (#6779)
Changes here allow the CPU compilation to be controlled using lowering
configuration that is specified on the operation. To plumb this through,
the lowering config can be used to specify a pass-pipeline and
workgroup size as well.
diff --git a/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp b/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp
index 9133eca..0abb16b 100644
--- a/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp
+++ b/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp
@@ -183,7 +183,6 @@
/// Sets the lowering configuration for dispatch region with root op being a
/// generic op.
static LogicalResult setDefaultRootConfig(FuncOp entryPointFn, Operation *op) {
- if (getLoweringConfig(op)) return success();
auto partitionedLoops = getPartitionedLoops(op);
if (partitionedLoops.empty()) {
// Return success without doing anything. Eventually default will be used.
@@ -213,15 +212,42 @@
for (auto computeOp : computeOps) {
if (!hasMarker(computeOp, getWorkgroupMarker())) continue;
- auto setRootConfigFn = [&](Operation *op) -> LogicalResult {
- return TypeSwitch<Operation *, LogicalResult>(op)
- .Case<linalg::Mmt4DOp, linalg::ContractionOpInterface>(
- [&](auto op) { return setRootConfig(entryPointFn, op); })
- .Default([&](Operation *op) { return success(); });
- };
-
- if (failed(setRootConfigFn(computeOp))) {
- return failure();
+ /// If the op already has a lowering config, then check for whether it
+ /// specifies a pass-pipeline and workgroup size as well. If so use those.
+ if (auto config = getLoweringConfig(computeOp)) {
+ IREE::HAL::DispatchLoweringPassPipeline passPipeline =
+ IREE::HAL::DispatchLoweringPassPipeline::CPUDefault;
+ if (auto passPipelineAttr = config.passPipeline()) {
+ passPipeline = passPipelineAttr.getValue();
+ }
+ SmallVector<int64_t, 4> workgroupSize;
+ if (auto workgroupSizeAttr = config.workgroupSize()) {
+ workgroupSize = llvm::to_vector<4>(
+ llvm::map_range(workgroupSizeAttr, [](Attribute intAttr) {
+ return intAttr.cast<IntegerAttr>().getInt();
+ }));
+ }
+ if (failed(setOpConfigAndEntryPointFnTranslation(
+ entryPointFn, computeOp, config, passPipeline, workgroupSize))) {
+ return failure();
+ }
+ // Reset the op configuration to drop the pass-pipeline and workgroup size
+ // info. The op does not carry that information anymore.
+ auto resetConfig = IREE::HAL::LoweringConfig::get(
+ config.tileSizes(), config.nativeVectorSize(),
+ /*passPipeline =*/nullptr,
+ /*workgroupSize =*/nullptr, computeOp->getContext());
+ setLoweringConfig(computeOp, resetConfig);
+ } else {
+ auto setRootConfigFn = [&](Operation *op) -> LogicalResult {
+ return TypeSwitch<Operation *, LogicalResult>(op)
+ .Case<linalg::Mmt4DOp, linalg::ContractionOpInterface>(
+ [&](auto op) { return setRootConfig(entryPointFn, op); })
+ .Default([&](Operation *op) { return success(); });
+ };
+ if (failed(setRootConfigFn(computeOp))) {
+ return failure();
+ }
}
if (getLoweringConfig(computeOp)) {
diff --git a/iree/compiler/Codegen/LLVMCPU/LLVMCPULowerExecutableTarget.cpp b/iree/compiler/Codegen/LLVMCPU/LLVMCPULowerExecutableTarget.cpp
index 85f47b8..cdb4720 100644
--- a/iree/compiler/Codegen/LLVMCPU/LLVMCPULowerExecutableTarget.cpp
+++ b/iree/compiler/Codegen/LLVMCPU/LLVMCPULowerExecutableTarget.cpp
@@ -150,6 +150,7 @@
executableLoweringPipeline.nest<ModuleOp>();
switch (passPipeline.getValue()) {
case IREE::HAL::DispatchLoweringPassPipeline::CPUDefault:
+ case IREE::HAL::DispatchLoweringPassPipeline::None:
addCPUDefaultPassPipeline(nestedModulePM);
break;
case IREE::HAL::DispatchLoweringPassPipeline::CPUVectorization:
diff --git a/iree/compiler/Codegen/LLVMCPU/test/materialize_launch_configuration.mlir b/iree/compiler/Codegen/LLVMCPU/test/materialize_launch_configuration.mlir
index 9c8a4c1..653f7b1 100644
--- a/iree/compiler/Codegen/LLVMCPU/test/materialize_launch_configuration.mlir
+++ b/iree/compiler/Codegen/LLVMCPU/test/materialize_launch_configuration.mlir
@@ -67,7 +67,11 @@
// -----
-// CHECK-NOT: #config
+// CHECK: #[[CONFIG:.+]] = {passPipeline = 0 : i32}
+// CHECK-NOT: #config
+// CHECK: hal.executable.entry_point @add_no_config
+// CHECK-SAME: translation.info = #[[CONFIG]]
+// CHECK-NOT: #config
hal.executable @add_no_config attributes {sym_visibility = "private"} {
hal.interface @io {
@@ -352,3 +356,80 @@
// CHECK: hal.return %[[D0]], %[[D1]], %[[ARG2]]
// CHECK: linalg.batch_matmul
// CHECK-SAME: lowering.config = #[[CONFIG]]
+
+// -----
+
+hal.executable @preset_config_matmul_tensors attributes {sym_visibility = "private"} {
+ hal.executable.variant @system_elf_x86_64, target = #hal.executable.target<"llvm", "system-elf-x86_64"> {
+ hal.executable.entry_point @preset_config attributes {interface = @io, ordinal = 0 : index}
+ builtin.module {
+ builtin.func @preset_config() {
+ %c0 = constant 0 : index
+ %c512 = constant 512 : index
+ %c128 = constant 128 : index
+ %cst = constant 0.000000e+00 : f32
+ %0 = hal.interface.binding.subspan @io::@s0b0_ro_external[%c0] : !flow.dispatch.tensor<readonly:128x256xf32>
+ %1 = hal.interface.binding.subspan @io::@s0b1_ro_external[%c0] : !flow.dispatch.tensor<readonly:256x512xf32>
+ %2 = hal.interface.binding.subspan @io::@s0b2_xw_external[%c0] : !flow.dispatch.tensor<writeonly:128x512xf32>
+ %workgroup_size_x = hal.interface.workgroup.size[0] : index
+ %workgroup_size_y = hal.interface.workgroup.size[1] : index
+ %workgroup_id_x = hal.interface.workgroup.id[0] : index
+ %workgroup_count_x = hal.interface.workgroup.count[0] : index
+ %workgroup_id_y = hal.interface.workgroup.id[1] : index
+ %workgroup_count_y = hal.interface.workgroup.count[1] : index
+ %3 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_id_y, %workgroup_size_y]
+ %4 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_count_y, %workgroup_size_y]
+ scf.for %arg0 = %3 to %c128 step %4 {
+ %5 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_id_x, %workgroup_size_x]
+ %6 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_count_x, %workgroup_size_x]
+ scf.for %arg1 = %5 to %c512 step %6 {
+ %7 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 128)>(%arg0)[%workgroup_size_y]
+ %8 = flow.dispatch.tensor.load %0, offsets = [%arg0, 0], sizes = [%7, 256], strides = [1, 1] : !flow.dispatch.tensor<readonly:128x256xf32> -> tensor<?x256xf32>
+ %9 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 512)>(%arg1)[%workgroup_size_x]
+ %10 = flow.dispatch.tensor.load %1, offsets = [0, %arg1], sizes = [256, %9], strides = [1, 1] : !flow.dispatch.tensor<readonly:256x512xf32> -> tensor<256x?xf32>
+ %11 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 128)>(%arg0)[%workgroup_size_y]
+ %12 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 512)>(%arg1)[%workgroup_size_x]
+ %13 = affine.min affine_map<(d0)[s0] -> (-d0 + 128, s0)>(%arg0)[%workgroup_size_y]
+ %14 = affine.min affine_map<(d0)[s0] -> (-d0 + 512, s0)>(%arg1)[%workgroup_size_x]
+ %15 = linalg.init_tensor [%13, %14] : tensor<?x?xf32>
+ %16 = linalg.fill(%cst, %15) : f32, tensor<?x?xf32> -> tensor<?x?xf32>
+ %17 = linalg.matmul {__internal_linalg_transform__ = "workgroup", lowering.config = {passPipeline = 1 : i32, tileSizes = [[32, 32, 32]]}} ins(%8, %10 : tensor<?x256xf32>, tensor<256x?xf32>) outs(%16 : tensor<?x?xf32>) -> tensor<?x?xf32>
+ flow.dispatch.tensor.store %17, %2, offsets = [%arg0, %arg1], sizes = [%11, %12], strides = [1, 1] : tensor<?x?xf32> -> !flow.dispatch.tensor<writeonly:128x512xf32>
+ }
+ }
+ return
+ }
+ hal.interface @io attributes {sym_visibility = "private"} {
+ hal.interface.binding @s0b0_ro_external, set=0, binding=0, type="StorageBuffer", access="Read"
+ hal.interface.binding @s0b1_ro_external, set=0, binding=1, type="StorageBuffer", access="Read"
+ hal.interface.binding @s0b2_xw_external, set=0, binding=2, type="StorageBuffer", access="Write|Discard"
+ }
+ }
+ }
+}
+// CHECK-DAG: #[[CONFIG:.+]] = {nativeVectorSize = [], tileSizes = {{\[}}[32, 32, 32]{{\]}}}
+// CHECK-DAG: #[[MAP0:.+]] = affine_map<()[s0] -> (s0 ceildiv 32)>
+// CHECK-DAG: #[[MAP1:.+]] = affine_map<()[s0] -> (s0 * 32)>
+// CHECK: hal.executable.entry_point
+// CHECK-SAME: translation.info = {passPipeline = 1 : i32, workloadPerWorkgroup = [32, 32]}
+// CHECK-NEXT: ^bb0(%[[ARG0:[a-zA-Z0-9]+]]: index, %[[ARG1:[a-zA-Z0-9]+]]: index
+// CHECK-DAG: %[[C1:.+]] = constant 1 : index
+// CHECK-DAG: %[[NWG_X:.+]] = affine.apply #[[MAP0]]()[%[[ARG0]]]
+// CHECK-DAG: %[[NWG_Y:.+]] = affine.apply #[[MAP0]]()[%[[ARG1]]]
+// CHECK: return %[[NWG_X]], %[[NWG_Y]], %[[C1]]
+// CHECK: builtin.module
+// CHECK: builtin.func @preset_config
+// CHECK-DAG: %[[WGID_X:.+]] = hal.interface.workgroup.id[0]
+// CHECK-DAG: %[[WGCOUNT_X:.+]] = hal.interface.workgroup.count[0]
+// CHECK-DAG: %[[WGID_Y:.+]] = hal.interface.workgroup.id[1]
+// CHECK-DAG: %[[WGCOUNT_Y:.+]] = hal.interface.workgroup.count[1]
+// CHECK: %[[LB_Y:.+]] = affine.apply #[[MAP1]]()[%[[WGID_Y]]]
+// CHECK: %[[STEP_Y:.+]] = affine.apply #[[MAP1]]()[%[[WGCOUNT_Y]]]
+// CHECK: scf.for %[[IV0:.+]] = %[[LB_Y]] to %{{.+}} step %[[STEP_Y]]
+// CHECK: %[[LB_X:.+]] = affine.apply #[[MAP1]]()[%[[WGID_X]]]
+// CHECK: %[[STEP_X:.+]] = affine.apply #[[MAP1]]()[%[[WGCOUNT_X]]]
+// CHECK: scf.for %[[IV1:.+]] = %[[LB_X]] to %{{.+}} step %[[STEP_X]]
+// CHECK: linalg.matmul
+// CHECK-SAME: lowering.config = #[[CONFIG]]
+// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<32x256xf32>, tensor<256x32xf32>)
+// CHECK-SAME: outs(%{{.+}} : tensor<32x32xf32>)
diff --git a/iree/compiler/Codegen/LLVMGPU/test/gpu_set_num_workgroups.mlir b/iree/compiler/Codegen/LLVMGPU/test/gpu_set_num_workgroups.mlir
index c5d0979..7034e05 100644
--- a/iree/compiler/Codegen/LLVMGPU/test/gpu_set_num_workgroups.mlir
+++ b/iree/compiler/Codegen/LLVMGPU/test/gpu_set_num_workgroups.mlir
@@ -143,14 +143,15 @@
}
}
-// CHECK-DAG: #[[CONFIG:.+]] = {tileSizes = {{\[}}[]{{\]}}}
+// CHECK-DAG: #[[CONFIG0:.+]] = {passPipeline = 2 : i32}
+// CHECK-DAG: #[[CONFIG1:.+]] = {tileSizes = {{\[}}[]{{\]}}}
// CHECK: hal.executable.entry_point @predict_dispatch_153
-// CHECK-SAME: passPipeline = 2 : i32
+// CHECK-SAME: translation.info = #[[CONFIG0]]
// CHECK-SAME: workgroup_size = [1 : index, 1 : index, 1 : index]
// CHECK-NEXT: ^bb0(%[[ARG0:[a-zA-Z0-9]+]]: index,
// CHECK-DAG: %[[C1:.+]] = constant 1 : index
// CHECK: hal.return %[[C1]], %[[C1]], %[[C1]]
// CHECK: linalg.fill
-// CHECK-SAME: lowering.config = #[[CONFIG]]
+// CHECK-SAME: lowering.config = #[[CONFIG1]]
// CHECK: linalg.generic
-// CHECK-SAME: lowering.config = #[[CONFIG]]
+// CHECK-SAME: lowering.config = #[[CONFIG1]]
diff --git a/iree/compiler/Codegen/Utils/Utils.cpp b/iree/compiler/Codegen/Utils/Utils.cpp
index 28a28b2..18df72f 100644
--- a/iree/compiler/Codegen/Utils/Utils.cpp
+++ b/iree/compiler/Codegen/Utils/Utils.cpp
@@ -81,31 +81,27 @@
}
LogicalResult setOpConfigAndEntryPointFnTranslation(
- FuncOp entryPointFn, Operation *op, TileSizesListTypeRef tileSizes,
- ArrayRef<int64_t> nativeVectorSizes,
+ FuncOp entryPointFn, Operation *op, IREE::HAL::LoweringConfig config,
IREE::HAL::DispatchLoweringPassPipeline passPipeline,
ArrayRef<int64_t> workgroupSize) {
- IREE::HAL::LoweringConfig config =
- buildConfigAttr(tileSizes, nativeVectorSizes, op->getContext());
- setLoweringConfig(op, config);
auto partitionedLoops = getPartitionedLoops(op);
SmallVector<int64_t, 3> workloadPerWorkgroup;
- if (!tileSizes.empty() && !tileSizes[0].empty() &&
- !partitionedLoops.empty()) {
+ auto tileSizes = getTileSizes(config, 0);
+ if (!tileSizes.empty() && !partitionedLoops.empty()) {
for (unsigned depth : partitionedLoops) {
- if (depth >= tileSizes[0].size()) {
+ if (depth >= tileSizes.size()) {
return op->emitOpError(
"illegal configuration for lowering op, expect first level "
"tile size to contain at least ")
<< partitionedLoops.back() << " elements";
}
- if (tileSizes[0][depth] == 0) {
+ if (tileSizes[depth] == 0) {
return op->emitOpError("illegal to set tilesize of loop ")
<< depth
<< " to zero since it is set to be partitioned at the flow "
"level";
}
- workloadPerWorkgroup.push_back(tileSizes[0][depth]);
+ workloadPerWorkgroup.push_back(tileSizes[depth]);
}
if (!workloadPerWorkgroup.empty()) {
workloadPerWorkgroup =
diff --git a/iree/compiler/Codegen/Utils/Utils.h b/iree/compiler/Codegen/Utils/Utils.h
index 6d195dd..cc7376a 100644
--- a/iree/compiler/Codegen/Utils/Utils.h
+++ b/iree/compiler/Codegen/Utils/Utils.h
@@ -43,20 +43,22 @@
/// formation to tile and distribute the ops.
SmallVector<unsigned> getPartitionedLoops(Operation *op);
-// /// Usually the tile sizes for the first level of tiling decides the
-// workgroup
-// /// size for the dispatch on the CPU backend. This is a general helper that
-// /// converts tile sizes of the first level into workgroup sizes.
-// SmallVector<int64_t, 3> getWorkloadPerWorkgroup(
-// ArrayRef<int64_t> firstLevelTileSizes, ArrayRef<int64_t>
-// partitionedLoops);
-
/// Sets translation for the entry-point function based on op configuration.
LogicalResult setOpConfigAndEntryPointFnTranslation(
- FuncOp entryPointFn, Operation *op, TileSizesListTypeRef tileSizes,
- ArrayRef<int64_t> nativeVectorSizes,
+ FuncOp entryPointFn, Operation *op, IREE::HAL::LoweringConfig config,
IREE::HAL::DispatchLoweringPassPipeline passPipeline,
ArrayRef<int64_t> workgroupSize = {});
+inline LogicalResult setOpConfigAndEntryPointFnTranslation(
+ FuncOp entryPointFn, Operation *op, TileSizesListTypeRef tileSizes,
+ ArrayRef<int64_t> nativeVectorSize,
+ IREE::HAL::DispatchLoweringPassPipeline passPipeline,
+ ArrayRef<int64_t> workgroupSize = {}) {
+ IREE::HAL::LoweringConfig config =
+ buildConfigAttr(tileSizes, nativeVectorSize, op->getContext());
+ setLoweringConfig(op, config);
+ return setOpConfigAndEntryPointFnTranslation(entryPointFn, op, config,
+ passPipeline, workgroupSize);
+}
/// Returns the number of outer parallel loops of a linalgOp.
/// Note: To be used only if needed. Use the `getPartitionedLoops` method if
diff --git a/iree/compiler/Dialect/HAL/IR/LoweringConfig.cpp b/iree/compiler/Dialect/HAL/IR/LoweringConfig.cpp
index 1c82197..1e1c2b5 100644
--- a/iree/compiler/Dialect/HAL/IR/LoweringConfig.cpp
+++ b/iree/compiler/Dialect/HAL/IR/LoweringConfig.cpp
@@ -109,7 +109,8 @@
nativeVectorSizeAttr = builder.getI64ArrayAttr(nativeVectorSize);
}
return IREE::HAL::LoweringConfig::get(tileSizesAttr, nativeVectorSizeAttr,
- context);
+ /*passPipeline = */ nullptr,
+ /*workgroupSize = */ nullptr, context);
}
TileSizesListType getTileSizes(IREE::HAL::LoweringConfig config) {
diff --git a/iree/compiler/Dialect/HAL/IR/LoweringConfig.td b/iree/compiler/Dialect/HAL/IR/LoweringConfig.td
index 519c754..f067014 100644
--- a/iree/compiler/Dialect/HAL/IR/LoweringConfig.td
+++ b/iree/compiler/Dialect/HAL/IR/LoweringConfig.td
@@ -26,6 +26,8 @@
: I32EnumAttrCase<"SPIRVVectorize", 6>;
def SPIRV_DistributeToGlobalID
: I32EnumAttrCase<"SPIRVDistributeToGlobalID", 7>;
+def None
+ : I32EnumAttrCase<"None", 0x7FFFFFFF>;
// EnumAttrCase for all known lowerings for ops within dispatch region
// to scalar/native-vector code.
@@ -33,8 +35,8 @@
"DispatchLoweringPassPipeline",
"identifier for pass pipeline use to lower dispatch region",
[CPU_Default, CPU_Vectorization, LLVMGPU_SimpleDistribute,
- LLVMGPU_Vectorize, LLVMGPU_MatmulSimt, SPIRV_SimpleDistribute, SPIRV_Vectorize,
- SPIRV_DistributeToGlobalID]> {
+ LLVMGPU_Vectorize, LLVMGPU_MatmulSimt, SPIRV_SimpleDistribute,
+ SPIRV_Vectorize, SPIRV_DistributeToGlobalID, None]> {
let cppNamespace = "::mlir::iree_compiler::IREE::HAL";
}
@@ -58,6 +60,12 @@
DefaultValuedAttr<TileSizesListAttr, "{}">>,
StructFieldAttr<"nativeVectorSize",
DefaultValuedAttr<I64ArrayAttr, "{}">>,
+ StructFieldAttr<"passPipeline",
+ DefaultValuedAttr<
+ DispatchLoweringPassPipelineEnum,
+ "IREE::HAL::DispatchLoweringPassPipeline::None">>,
+ StructFieldAttr<"workgroupSize",
+ DefaultValuedAttr<I64ArrayAttr, "{}">>
]>;
#endif // IREE_COMPILER_DIALECT_HAL_IR_LOWERINGCONFIG
diff --git a/iree/test/e2e/regression/BUILD b/iree/test/e2e/regression/BUILD
index 0946c8e..b95db69 100644
--- a/iree/test/e2e/regression/BUILD
+++ b/iree/test/e2e/regression/BUILD
@@ -50,6 +50,7 @@
"dynamic_linalg_matmul_on_tensors_fuse_1.mlir",
"dynamic_linalg_matmul_on_tensors_fuse_2.mlir",
"linalg_ext_ops.mlir",
+ "lowering_config.mlir",
] + BACKEND_TESTS,
),
data = [
@@ -65,6 +66,7 @@
name = "check_regression_linalg_ops_dylib-llvm-aot",
srcs = [
"linalg_ext_ops.mlir",
+ "lowering_config.mlir",
] + BACKEND_TESTS,
compiler_flags = ["-iree-input-type=mhlo"],
driver = "dylib",
diff --git a/iree/test/e2e/regression/CMakeLists.txt b/iree/test/e2e/regression/CMakeLists.txt
index 3e6bc37..38b2886 100644
--- a/iree/test/e2e/regression/CMakeLists.txt
+++ b/iree/test/e2e/regression/CMakeLists.txt
@@ -41,6 +41,7 @@
SRCS
"linalg_ext_ops.mlir"
"linalg_ops.mlir"
+ "lowering_config.mlir"
TARGET_BACKEND
"dylib-llvm-aot"
DRIVER
diff --git a/iree/test/e2e/regression/lowering_config.mlir b/iree/test/e2e/regression/lowering_config.mlir
new file mode 100644
index 0000000..a70c491
--- /dev/null
+++ b/iree/test/e2e/regression/lowering_config.mlir
@@ -0,0 +1,12 @@
+#config1 = {tileSizes = [[32, 32, 32]], passPipeline = 1 : i32}
+#config2 = {tileSizes = [[64, 64, 64]], passPipeline = 1 : i32}
+func @lowering_config_test() {
+ %a = util.unfoldable_constant dense<1.0> : tensor<128x256xf32>
+ %b = util.unfoldable_constant dense<2.0> : tensor<256x512xf32>
+ %c = util.unfoldable_constant dense<2.0> : tensor<256x1024xf32>
+ %d = "mhlo.dot"(%a, %b) {lowering.config = #config1} : (tensor<128x256xf32>, tensor<256x512xf32>) -> tensor<128x512xf32>
+ %e = "mhlo.dot"(%a, %c) {lowering.config = #config2} : (tensor<128x256xf32>, tensor<256x1024xf32>) -> tensor<128x1024xf32>
+ check.expect_almost_eq_const(%d, dense<512.0> : tensor<128x512xf32>) : tensor<128x512xf32>
+ check.expect_almost_eq_const(%e, dense<512.0> : tensor<128x1024xf32>) : tensor<128x1024xf32>
+ return
+}
