[Codegen] KernelDispatch recognition and tiling for data-tiled conv (#24716)
Teaches LLVMCPU KernelDispatch to recognize data-tiled convolution
dispatches produced by encoding materialization and assigns them a
dedicated root lowering config that vectorizes the output width (ow) and
the inner data-tiling channel dims (oc_inner, ic_inner, sizes k0/c0),
and distributes batch (unit tiles), output channel blocks (oc_outer),
and output height, with peeling enabled.
Co-authored-by: Jelle Schuhmacher <schuehmacher@roofline.ai>
Assisted-by: Claude code
Signed-off-by: Pooja Hemashekar <hemashekar@roofline.ai>
---------
Signed-off-by: Pooja Hemashekar <hemashekar@roofline.ai>
Co-authored-by: Jelle Schuhmacher <schuehmacher@roofline.ai>
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.cpp b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.cpp
index d49aad4..12704af 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.cpp
+++ b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.cpp
@@ -10,6 +10,8 @@
#include "iree/compiler/Dialect/LinalgExt/Utils/MatchUtils.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/InterleavedRange.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Linalg/IR/LinalgInterfaces.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
@@ -690,4 +692,63 @@
return space;
}
+/// Recognize the data-tiled convolution generic emitted by encoding
+/// materialization, identified by its iterator types [P,P,P,P, R,R,R, P,R],
+/// operand ranks (input 5, filter 6, output 5), indexing maps, and
+/// multiply-accumulate body.
+bool isDataTiledConvGeneric(Operation *op) {
+ auto genericOp = dyn_cast<linalg::GenericOp>(op);
+ if (!genericOp) {
+ return false;
+ }
+ using IT = utils::IteratorType;
+ SmallVector<IT> iterTypes = genericOp.getIteratorTypesArray();
+ SmallVector<IT> expected = {IT::parallel, IT::parallel, IT::parallel,
+ IT::parallel, IT::reduction, IT::reduction,
+ IT::reduction, IT::parallel, IT::reduction};
+ if (iterTypes.size() != 9 || !llvm::equal(iterTypes, expected)) {
+ return false;
+ }
+
+ if (genericOp.getNumDpsInputs() != 2 || genericOp.getNumDpsInits() != 1) {
+ return false;
+ }
+ auto inputType = dyn_cast<RankedTensorType>(
+ genericOp.getDpsInputOperand(0)->get().getType());
+ auto filterType = dyn_cast<RankedTensorType>(
+ genericOp.getDpsInputOperand(1)->get().getType());
+ auto outputType = dyn_cast<RankedTensorType>(
+ genericOp.getDpsInitOperand(0)->get().getType());
+ if (!inputType || !filterType || !outputType) {
+ return false;
+ }
+ if (inputType.getRank() != 5 || filterType.getRank() != 6 ||
+ outputType.getRank() != 5) {
+ return false;
+ }
+
+ FailureOr<linalg::ConvolutionDimensions> cDims =
+ linalg::inferConvolutionDims(genericOp);
+ if (failed(cDims) || cDims->strides.size() != 2 ||
+ cDims->dilations.size() != 2) {
+ return false;
+ }
+
+ // Require the maps and iterator types to be exactly what the encoding
+ // materialization emits for this op's strides/dilations.
+ DataTiledConvIterationSpace space = getDataTiledConvIterationSpace(
+ genericOp.getContext(), cDims->strides, cDims->dilations);
+ if (genericOp.getIndexingMapsArray() != space.indexingMaps ||
+ genericOp.getIteratorTypesArray() != space.iteratorTypes) {
+ return false;
+ }
+
+ return linalg::detail::isContractionBody(
+ *genericOp.getBlock(),
+ [](Operation *mul, Operation *add) {
+ return isa<arith::MulFOp>(mul) && isa<arith::AddFOp>(add);
+ },
+ llvm::nulls());
+}
+
} // namespace mlir::iree_compiler::IREE::Codegen
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.h b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.h
index b9c509d..5bfafb6 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.h
+++ b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.h
@@ -165,6 +165,11 @@
getDataTiledConvIterationSpace(MLIRContext *ctx, ArrayRef<int64_t> strides,
ArrayRef<int64_t> dilations);
+/// Returns true if `op` is a 9D data-tiled convolution generic emitted by
+/// encoding materialization (matched by iterator types, operand ranks,
+/// indexing maps, and multiply-accumulate body).
+bool isDataTiledConvGeneric(Operation *op);
+
} // namespace mlir::iree_compiler::IREE::Codegen
#endif // IREE_COMPILER_CODEGEN_DIALECT_CODEGEN_UTILS_H_
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/BUILD.bazel b/compiler/src/iree/compiler/Codegen/LLVMCPU/BUILD.bazel
index c1e522d..4e81511 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/BUILD.bazel
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/BUILD.bazel
@@ -102,6 +102,7 @@
"//compiler/src/iree/compiler/Codegen/Dialect/CPU/IR:IREECPUDialect",
"//compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR:IREECodegenDialect",
"//compiler/src/iree/compiler/Codegen/Dialect/Codegen/Transforms:IREECodegenTransforms",
+ "//compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils",
"//compiler/src/iree/compiler/Codegen/Dialect/VectorExt/IR:IREEVectorExtDialect",
"//compiler/src/iree/compiler/Codegen/Dialect/VectorExt/Transforms:VectorExtTransforms",
"//compiler/src/iree/compiler/Codegen/Interfaces:PartitionableLoopsInterface",
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/LLVMCPU/CMakeLists.txt
index 1a5de09..d20f744 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/CMakeLists.txt
@@ -160,6 +160,7 @@
iree::compiler::Codegen::Dialect::CPU::IR::IREECPUDialect
iree::compiler::Codegen::Dialect::Codegen::IR::IREECodegenDialect
iree::compiler::Codegen::Dialect::Codegen::Transforms::IREECodegenTransforms
+ iree::compiler::Codegen::Dialect::Codegen::Utils
iree::compiler::Codegen::Dialect::VectorExt::IR::IREEVectorExtDialect
iree::compiler::Codegen::Dialect::VectorExt::Transforms::VectorExtTransforms
iree::compiler::Codegen::Interfaces::PartitionableLoopsInterface
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp b/compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp
index ede7559..83c8444 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp
@@ -12,6 +12,7 @@
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenInterfaces.h"
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenOps.h"
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenTypes.h"
+#include "iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.h"
#include "iree/compiler/Codegen/Interfaces/PartitionableLoopsInterface.h"
#include "iree/compiler/Codegen/LLVMCPU/LLVMCPUSelectUKernels.h"
#include "iree/compiler/Codegen/LLVMCPU/TargetMLTransformInfo.h"
@@ -3114,6 +3115,65 @@
getCPUTranslationInfo(op.getContext(), CPUPipeline::Mmt4dTilingExpert));
}
+/// Assigns CPUDoubleTilingExpert pipeline config to the data-tiled
+/// convolution generic.
+static LogicalResult
+setConvDataTiledGenericRootConfig(mlir::FunctionOpInterface entryPointFn,
+ linalg::GenericOp convOp) {
+ // Loop ranges:
+ // (d0, d1, d2, d3, d4, d5, d6, d7, d8)
+ // (n, OC/k0, OH, OW, IC/c0, FH, FW, k0, c0)
+ SmallVector<int64_t, 9> loopRanges =
+ cast<linalg::LinalgOp>(convOp.getOperation()).getStaticLoopRanges();
+ int64_t k0 = loopRanges[7];
+ int64_t c0 = loopRanges[8];
+
+ // Vectorization strategy:
+ // - Vectorize across OW, the register-blocking dimension: the kernel
+ // carries one accumulator vector per OW element to hide FMA latency.
+ // Use k0 as the OW tile size for now; it matches the desired
+ // accumulator count on current targets (16 for AVX-512, 8 for NEON).
+ // TODO(phemashekar): Derive this directly from the target register budget
+ // instead.
+ // - Vectorize across the data-tiling inner dims oc_inner (k0) and
+ // ic_inner (c0).
+ SmallVector<int64_t> vecTileSizes(9, 1);
+ vecTileSizes[3] = k0;
+ vecTileSizes[7] = k0;
+ vecTileSizes[8] = c0;
+ setAlwaysVectorizeSizes(convOp, vecTileSizes);
+
+ // Distribute over N, OC/k0, and OH. N is limited to unit tiles: batch
+ // tile sizes are constrained by the generated temporary buffer sizes,
+ // and larger batch tiles can lead to stack allocation errors.
+ DistributionHeuristicConfig distConfig;
+ distConfig.maxTileSizes.assign(9, 0);
+ distConfig.maxTileSizes[0] = 1;
+ distConfig.maxTileSizes[1] = clDefaultDistTileSize / 2;
+ distConfig.maxTileSizes[2] = clDefaultDistTileSize / 2;
+ SmallVector<int64_t> distTileSizes =
+ getDefaultDistributedLevelTileSizes(convOp, distConfig);
+
+ LDBG() << "Data tiled convolution:";
+ LDBG() << " Dist tile sizes: " << distTileSizes;
+ LDBG() << " Vector tile sizes: " << vecTileSizes;
+
+ LoweringConfigGenerator generator(convOp);
+ generator.setDistributionTileSizes(distTileSizes);
+ generator.setVectorTileSizes(vecTileSizes);
+ IREE::CPU::LoweringConfigAttr loweringConfig =
+ generator.generateCPULoweringConfig();
+
+ // Enable loop peeling so the OW vectorization tile produces static-sized
+ // main loop iterations (vectorizable) with a scalar peeled tail.
+ DictionaryAttr pipelineConfig =
+ getPipelineConfWithPeelingAttr(convOp.getContext());
+ return setOpConfigAndEntryPointFnTranslation(
+ entryPointFn, convOp, loweringConfig,
+ getCPUTranslationInfo(convOp.getContext(),
+ CPUPipeline::DoubleTilingExpert, pipelineConfig));
+}
+
/// Redirects to methods that set the configuration based on operation type.
static LogicalResult
setRootConfigImpl(mlir::FunctionOpInterface entryPointFn, Operation *op,
@@ -3147,6 +3207,10 @@
is2DPoolingOp(linalgOp)) {
return setConvInterfaceRootConfig(entryPointFn, linalgOp);
}
+ if (IREE::Codegen::isDataTiledConvGeneric(op)) {
+ return setConvDataTiledGenericRootConfig(entryPointFn,
+ cast<linalg::GenericOp>(op));
+ }
if (linalg::isaContractionOpInterface(linalgOp) &&
meetLegacyContractionOpInterface(linalgOp)) {
return setContractionRootConfig(entryPointFn, linalgOp);
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_aarch64_lowering_strategy.mlir b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_aarch64_lowering_strategy.mlir
index 56f07be..23dccd6 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_aarch64_lowering_strategy.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_aarch64_lowering_strategy.mlir
@@ -104,6 +104,28 @@
// -----
+#executable_target_embedded_elf_arm_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-arm_64", {data_layout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128", native_vector_size = 16 : index, target_triple = "aarch64-none-elf"}>
+#map = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d0, d4, d2 + d5, d3 + d6, d8)>
+#map1 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d1, d4, d5, d6, d8, d7)>
+#map2 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d0, d1, d2, d3, d7)>
+func.func @conv_2d_nchwc_data_tiled(%arg0: tensor<1x1x16x16x16xf32>, %arg1: tensor<1x1x3x3x16x16xf32>, %arg2: tensor<1x1x14x14x16xf32>) -> tensor<1x1x14x14x16xf32> attributes {hal.executable.target = #executable_target_embedded_elf_arm_64_} {
+ %0 = linalg.generic {indexing_maps = [#map, #map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction", "reduction", "parallel", "reduction"]} ins(%arg0, %arg1 : tensor<1x1x16x16x16xf32>, tensor<1x1x3x3x16x16xf32>) outs(%arg2 : tensor<1x1x14x14x16xf32>) {
+ ^bb0(%in: f32, %in_0: f32, %out: f32):
+ %1 = arith.mulf %in, %in_0 : f32
+ %2 = arith.addf %1, %out : f32
+ linalg.yield %2 : f32
+ } -> tensor<1x1x14x14x16xf32>
+ return %0 : tensor<1x1x14x14x16xf32>
+}
+// CHECK-DAG: #[[CONFIG:.+]] = #iree_cpu.lowering_config<distribution = [0, 0, 1, 0, 0, 0, 0, 0, 0], vector_common_parallel = [1, 1, 1, 16, 0, 0, 0, 16, 0], vector_reduction = [0, 0, 0, 0, 1, 1, 1, 0, 16]>
+// CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = #iree_cpu.pipeline<DoubleTilingExpert>, {enable_loop_peeling}>
+// CHECK: func.func @conv_2d_nchwc_data_tiled(
+// CHECK-SAME: translation_info = #[[TRANSLATION]]
+// CHECK: linalg.generic
+// CHECK-SAME: lowering_config = #[[CONFIG]]
+
+// -----
+
#executable_target_system_elf_arm_64_ = #hal.executable.target<"llvm-cpu", "system-elf-arm_64", {data_layout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128", native_vector_size = 16 : index, target_triple = "aarch64-none-linux-android30"}>
func.func @restrict_num_workgroups(%3: tensor<1x11x11x576xf32>, %4: tensor<5x5x576xf32>) -> tensor<1x7x7x576xf32> attributes {hal.executable.target = #executable_target_system_elf_arm_64_} {
%cst = arith.constant 0.000000e+00 : f32
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_x86_64_lowering_strategy.mlir b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_x86_64_lowering_strategy.mlir
index 3317074..8fc1c54 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_x86_64_lowering_strategy.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_x86_64_lowering_strategy.mlir
@@ -455,6 +455,30 @@
// -----
+// The materialized data-tiled convolution generic is recognized via the
+// convolution interface and routed to the double-tiling pipeline.
+#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu_features = "+avx512f", data_layout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128", native_vector_size = 16 : index, target_triple = "x86_64-unknown-linux-gnu"}>
+#map = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d0, d4, d2 + d5, d3 + d6, d8)>
+#map1 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d1, d4, d5, d6, d8, d7)>
+#map2 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d0, d1, d2, d3, d7)>
+func.func @conv_2d_nchwc_data_tiled(%arg0: tensor<1x1x16x16x16xf32>, %arg1: tensor<1x1x3x3x16x16xf32>, %arg2: tensor<1x1x14x14x16xf32>) -> tensor<1x1x14x14x16xf32> attributes {hal.executable.target = #executable_target_embedded_elf_x86_64_} {
+ %0 = linalg.generic {indexing_maps = [#map, #map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction", "reduction", "parallel", "reduction"]} ins(%arg0, %arg1 : tensor<1x1x16x16x16xf32>, tensor<1x1x3x3x16x16xf32>) outs(%arg2 : tensor<1x1x14x14x16xf32>) {
+ ^bb0(%in: f32, %in_0: f32, %out: f32):
+ %1 = arith.mulf %in, %in_0 : f32
+ %2 = arith.addf %1, %out : f32
+ linalg.yield %2 : f32
+ } -> tensor<1x1x14x14x16xf32>
+ return %0 : tensor<1x1x14x14x16xf32>
+}
+// CHECK-DAG: #[[CONFIG:.+]] = #iree_cpu.lowering_config<distribution = [0, 0, 1, 0, 0, 0, 0, 0, 0], vector_common_parallel = [1, 1, 1, 16, 0, 0, 0, 16, 0], vector_reduction = [0, 0, 0, 0, 1, 1, 1, 0, 16]>
+// CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = #iree_cpu.pipeline<DoubleTilingExpert>, {enable_loop_peeling}>
+// CHECK: func.func @conv_2d_nchwc_data_tiled(
+// CHECK-SAME: translation_info = #[[TRANSLATION]]
+// CHECK: linalg.generic
+// CHECK-SAME: lowering_config = #[[CONFIG]]
+
+// -----
+
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "cascadelake", cpu_features = "+mmx,+popcnt,+sse,+sse2,+sse3,+ssse3,+sse4.1,+sse4.2,+avx,+avx2,+fma,+avx512f,+bmi,+bmi2,+aes,+pclmul,+avx512vl,+avx512bw,+avx512dq,+avx512cd,+avx512vnni,+adx,+clflushopt,+clwb,+cx16,+cx8,+crc32,+f16c,+fsgsbase,+fxsr,+invpcid,+lzcnt,+movbe,+pku,+prfchw,+rdrnd,+rdseed,+sahf,+x87,+xsave,+xsavec,+xsaveopt,+xsaves", data_layout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128", native_vector_size = 32 : index, target_triple = "x86_64-none-elf", ukernels = false}>
func.func @pooling_nchw_max(%2: tensor<1x64x114x114xf32>) -> tensor<1x64x56x56xf32> attributes {hal.executable.target = #executable_target_embedded_elf_x86_64_} {
%cst = arith.constant -3.40282347E+38 : f32
