[Codegen] Lower data-tiled conv to conv_nchwc ukernel (#24730)
When the convolution ukernel is enabled, routes data-tiled convolution
dispatches through `Mmt4dTilingExpert` pipeline, where the 9D NCHWc
generic is recognized and lowered to an `iree_uk_conv_nchwc` microkernel
call.
Assisted-by: Claude code
---------
Signed-off-by: Pooja Hemashekar <hemashekar@roofline.ai>
diff --git a/compiler/src/iree/compiler/Codegen/Common/CPU/BUILD.bazel b/compiler/src/iree/compiler/Codegen/Common/CPU/BUILD.bazel
index 7f40cd7..1ef369d 100644
--- a/compiler/src/iree/compiler/Codegen/Common/CPU/BUILD.bazel
+++ b/compiler/src/iree/compiler/Codegen/Common/CPU/BUILD.bazel
@@ -84,6 +84,7 @@
"@llvm-project//mlir:FunctionInterfaces",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:LinalgDialect",
+ "@llvm-project//mlir:LinalgInterfaces",
"@llvm-project//mlir:LinalgTransforms",
"@llvm-project//mlir:LinalgUtils",
"@llvm-project//mlir:MemRefDialect",
diff --git a/compiler/src/iree/compiler/Codegen/Common/CPU/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/Common/CPU/CMakeLists.txt
index 0a52a37..c588002 100644
--- a/compiler/src/iree/compiler/Codegen/Common/CPU/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Codegen/Common/CPU/CMakeLists.txt
@@ -59,6 +59,7 @@
MLIRFunctionInterfaces
MLIRIR
MLIRLinalgDialect
+ MLIRLinalgInterfacesIncGenLib
MLIRLinalgTransforms
MLIRLinalgUtils
MLIRMemRefDialect
diff --git a/compiler/src/iree/compiler/Codegen/Common/CPU/CPULowerToUKernels.cpp b/compiler/src/iree/compiler/Codegen/Common/CPU/CPULowerToUKernels.cpp
index bfaa28e..fdecb81 100644
--- a/compiler/src/iree/compiler/Codegen/Common/CPU/CPULowerToUKernels.cpp
+++ b/compiler/src/iree/compiler/Codegen/Common/CPU/CPULowerToUKernels.cpp
@@ -11,6 +11,7 @@
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenDialect.h"
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenOps.h"
#include "iree/compiler/Codegen/Dialect/Codegen/IR/UKernelOps.h"
+#include "iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.h"
#include "iree/compiler/Codegen/Utils/Utils.h"
#include "iree/compiler/Dialect/Encoding/IR/EncodingOps.h"
#include "iree/compiler/Dialect/Encoding/IR/EncodingTypes.h"
@@ -18,6 +19,7 @@
#include "llvm/ADT/Repeated.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/Linalg/Utils/Utils.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/Attributes.h"
@@ -267,6 +269,121 @@
genericMicroKernelOp.getOperation());
}
+/// Matches a 9D data-tiled conv generic op and rewrites it to a call to
+/// iree_uk_conv_nchwc, that is later lowered into a call to the
+/// microkernel.
+static FailureOr<IREE::Codegen::UKernelOpInterface>
+matchDAGForUKernel(RewriterBase &rewriter, linalg::GenericOp op,
+ bool /*skipIntermediateRoundings*/) {
+ auto targetAttr = IREE::HAL::ExecutableTargetAttr::lookup(op);
+ const char ukernelName[] = "conv_nchwc";
+ if (!targetAttr || !hasUkernel(targetAttr.getConfiguration(), ukernelName)) {
+ return failure();
+ }
+ if (!IREE::Codegen::isDataTiledConvGeneric(op)) {
+ return rewriter.notifyMatchFailure(op, "not a data-tiled conv generic");
+ }
+ Value input = op.getDpsInputOperand(0)->get();
+ Value filter = op.getDpsInputOperand(1)->get();
+ Value output = op.getDpsInitOperand(0)->get();
+ auto inputType = cast<RankedTensorType>(input.getType());
+ auto filterType = cast<RankedTensorType>(filter.getType());
+ auto outputType = cast<RankedTensorType>(output.getType());
+
+ Type inElem = inputType.getElementType();
+ Type filterElem = filterType.getElementType();
+ Type outElem = outputType.getElementType();
+ uint32_t flags = 0;
+ if (inElem.isF32() && filterElem.isF32() && outElem.isF32()) {
+ flags = IREE_UK_FLAG_CONV_NCHWC_TYPE_F32F32F32;
+ } else {
+ return rewriter.notifyMatchFailure(op,
+ "unsupported conv_nchwc element types");
+ }
+
+ auto cDims = linalg::inferConvolutionDims(op);
+ if (failed(cDims) || cDims->strides.size() != 2 ||
+ cDims->dilations.size() != 2) {
+ return rewriter.notifyMatchFailure(op, "failed to infer conv dims");
+ }
+ // TODO(#24760): plumb dilation_h/w as ukernel operands and honor them
+ // in the ukernels (window expression currently assumes unit dilation).
+ if (cDims->dilations[0] != 1 || cDims->dilations[1] != 1) {
+ return rewriter.notifyMatchFailure(op, "only dilation=1 supported for now");
+ }
+ int64_t strideH = cDims->strides[0];
+ int64_t strideW = cDims->strides[1];
+
+ if (isInitializedToZero(output)) {
+ // Not setting flags |= IREE_UK_FLAG_CONV_NCHWC_ACCUMULATE, so the conv op
+ // won't read the existing accumulator, so its defining op can be discarded.
+ if (auto fillOp = output.getDefiningOp<linalg::FillOp>()) {
+ output = fillOp.getDpsInitOperand(0)->get();
+ }
+ } else {
+ // Tell the conv op to read the existing accumulator.
+ flags |= IREE_UK_FLAG_CONV_NCHWC_ACCUMULATE;
+ }
+
+ flags |= IREE_UK_FLAG_CONV_NCHWC_ALLOW_GENERIC_FALLBACK_TILE_FUNCTION;
+
+ Location loc = op.getLoc();
+
+ auto dimAsIndex = [&](Value v, int64_t d) -> Value {
+ return tensor::DimOp::create(rewriter, loc, v, d);
+ };
+
+ // isDataTiledConvGeneric() guarantees the canonical operand layouts.
+ Value n = dimAsIndex(input, 0);
+ Value icOuter = dimAsIndex(input, 1);
+ Value ocOuter = dimAsIndex(filter, 0);
+ Value fh = dimAsIndex(filter, 2);
+ Value fw = dimAsIndex(filter, 3);
+ Value k0 = arith::IndexCastOp::create(rewriter, loc, rewriter.getI32Type(),
+ dimAsIndex(filter, 5));
+ Value c0 = arith::IndexCastOp::create(rewriter, loc, rewriter.getI32Type(),
+ dimAsIndex(filter, 4));
+ Value oh = dimAsIndex(output, 2);
+ Value ow = dimAsIndex(output, 3);
+
+ Value strideHVal = arith::ConstantOp::create(
+ rewriter, loc, rewriter.getI32IntegerAttr(strideH));
+ Value strideWVal = arith::ConstantOp::create(
+ rewriter, loc, rewriter.getI32IntegerAttr(strideW));
+ Value flagsVal = arith::ConstantOp::create(rewriter, loc,
+ rewriter.getI32IntegerAttr(flags));
+
+ auto fn = getFnNameAndDefAttrs(ukernelName, rewriter, targetAttr);
+ SmallVector<Type> returnTypes =
+ getUKernelGenericReturnTypes(targetAttr, outputType);
+
+ // Ukernel signature order: (N, OC_outer, OH, OW, IC_outer, FH, FW, k0, c0,
+ // stride_h, stride_w, flags).
+ SmallVector<Value> otherOperands{n, ocOuter, oh, ow,
+ icOuter, fh, fw, k0,
+ c0, strideHVal, strideWVal, flagsVal};
+ // Per operand, the strides of the outer dims are passed to the ukernel. These
+ // operands are views of large packed tensors (via tensor.extract_slice), so
+ // each outer dim's stride comes from the parent layout and can't be assumed
+ // from the slice shape. We only block input/output channels, so the inner
+ // block dims (c0, k0) are contiguous and walked with compile-time strides and
+ // they don't require runtime strides; the remaining iterated dims need one.
+ // input [N, IC/c0, H, W, c0]: strides of dims 0, 1, 2 ->
+ // (input_stride_n, input_stride_ic_outer, input_stride_h)
+ // filter [OC/k0, IC/c0, FH, FW, c0, k0]: strides of dims 0, 1, 2 ->
+ // (filter_stride_oc_outer, filter_stride_ic_outer, filter_stride_fh)
+ // output [N, OC/k0, OH, OW, k0]: strides of dims 0, 1, 2 ->
+ // (output_stride_n, output_stride_oc_outer, output_stride_oh)
+ SmallVector<SmallVector<int64_t>> stridedDims = {
+ {0, 1, 2}, {0, 1, 2}, {0, 1, 2}};
+ auto genericMicroKernelOp = IREE::Codegen::UKernelGenericOp::create(
+ rewriter, loc, returnTypes, fn.name, ValueRange{input, filter}, output,
+ otherOperands,
+ /*fn_def_attrs=*/rewriter.getDictionaryAttr(fn.defAttrs), stridedDims);
+ return cast<IREE::Codegen::UKernelOpInterface>(
+ genericMicroKernelOp.getOperation());
+}
+
static FailureOr<IREE::Codegen::UKernelOpInterface>
matchDAGForUKernel(RewriterBase &rewriter, linalg::PackOp op,
bool /*skipIntermediateRoundings*/) {
@@ -653,7 +770,8 @@
auto allTargets = [](auto target) { return true; };
patterns.insert<LowerToUKernelPattern<linalg::Mmt4DOp>,
LowerToUKernelPattern<linalg::PackOp>,
- LowerToUKernelPattern<linalg::UnPackOp>>(
+ LowerToUKernelPattern<linalg::UnPackOp>,
+ LowerToUKernelPattern<linalg::GenericOp>>(
context, allTargets, skipIntermediateRoundings);
// These patterns are inherently specific to the VMVX backend.
patterns.insert<LowerToUKernelPattern<IREE::Codegen::QueryTileSizesOp>>(
diff --git a/compiler/src/iree/compiler/Codegen/Common/CPU/test/lower_to_ukernel_ops.mlir b/compiler/src/iree/compiler/Codegen/Common/CPU/test/lower_to_ukernel_ops.mlir
index a823864..ff1f9bc 100644
--- a/compiler/src/iree/compiler/Codegen/Common/CPU/test/lower_to_ukernel_ops.mlir
+++ b/compiler/src/iree/compiler/Codegen/Common/CPU/test/lower_to_ukernel_ops.mlir
@@ -597,3 +597,146 @@
// CHECK-SAME: ins(%[[ARG0]], %[[ARG1]] :
// CHECK-SAME: outs(%[[ARG2]] :
// CHECK: return %[[MICRO_KERNEL]]#0
+
+// -----
+
+#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_nchwc_f32f32f32(%input: tensor<1x1x16x16x16xf32>, %filter: tensor<1x1x3x3x16x16xf32>, %output: tensor<1x1x14x14x16xf32>) -> tensor<1x1x14x14x16xf32> attributes {
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz", {ukernels = "all", target_triple = "x86_64-xyz-xyz", cpu_features = "+avx512f"}>
+} {
+ %0 = linalg.generic {indexing_maps = [#map, #map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction", "reduction", "parallel", "reduction"]} ins(%input, %filter : tensor<1x1x16x16x16xf32>, tensor<1x1x3x3x16x16xf32>) outs(%output : 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-LABEL: func @conv_nchwc_f32f32f32(
+// CHECK-SAME: %[[INPUT:[a-zA-Z0-9]+]]: tensor<1x1x16x16x16xf32>
+// CHECK-SAME: %[[FILTER:[a-zA-Z0-9]+]]: tensor<1x1x3x3x16x16xf32>
+// CHECK-SAME: %[[OUTPUT:[a-zA-Z0-9]+]]: tensor<1x1x14x14x16xf32>
+// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
+// CHECK-DAG: %[[C3:.+]] = arith.constant 3 : index
+// CHECK-DAG: %[[C14:.+]] = arith.constant 14 : index
+// CHECK-DAG: %[[C16_I32:.+]] = arith.constant 16 : i32
+// CHECK-DAG: %[[C1_I32:.+]] = arith.constant 1 : i32
+// Flag constant: 769 = 0x301 = ACCUMULATE (0x100) | ALLOW_GENERIC_FALLBACK_TILE_FUNCTION (0x200) | TYPE_F32F32F32 (0x1)
+// CHECK-DAG: %[[FLAGS:.+]] = arith.constant 769 : i32
+// CHECK: %[[MICRO_KERNEL:.+]]:2 = iree_codegen.ukernel.generic "iree_uk_conv_nchwc"
+// CHECK-SAME: ins(%[[INPUT]], %[[FILTER]] :
+// CHECK-SAME: outs(%[[OUTPUT]] :
+// CHECK-SAME: (%[[C1]], %[[C1]], %[[C14]], %[[C14]], %[[C1]], %[[C3]], %[[C3]], %[[C16_I32]], %[[C16_I32]], %[[C1_I32]], %[[C1_I32]], %[[FLAGS]] :
+// CHECK-SAME: strided_dims({{\[}}[0, 1, 2], [0, 1, 2], [0, 1, 2]])
+// CHECK: return %[[MICRO_KERNEL]]#0
+
+// -----
+
+// Zero-initialized accumulator: the linalg.fill is folded away and the ukernel
+// writes a fresh buffer.
+
+#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_nchwc_zero_fill(%input: tensor<1x1x16x16x16xf32>, %filter: tensor<1x1x3x3x16x16xf32>) -> tensor<1x1x14x14x16xf32> attributes {
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz", {ukernels = "all", target_triple = "x86_64-xyz-xyz", cpu_features = "+avx512f"}>
+} {
+ %cst = arith.constant 0.000000e+00 : f32
+ %empty = tensor.empty() : tensor<1x1x14x14x16xf32>
+ %fill = linalg.fill ins(%cst : f32) outs(%empty : tensor<1x1x14x14x16xf32>) -> tensor<1x1x14x14x16xf32>
+ %0 = linalg.generic {indexing_maps = [#map, #map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction", "reduction", "parallel", "reduction"]} ins(%input, %filter : tensor<1x1x16x16x16xf32>, tensor<1x1x3x3x16x16xf32>) outs(%fill : 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-LABEL: func @conv_nchwc_zero_fill(
+// Flag constant: 513 = 0x201 = ALLOW_GENERIC_FALLBACK_TILE_FUNCTION (0x200) | TYPE_F32F32F32 (0x1); no ACCUMULATE since the fill zeroes the accumulator
+// CHECK-DAG: %[[FLAGS:.+]] = arith.constant 513 : i32
+// CHECK-DAG: %[[EMPTY:.+]] = tensor.empty() : tensor<1x1x14x14x16xf32>
+// CHECK: %[[MICRO_KERNEL:.+]]:2 = iree_codegen.ukernel.generic "iree_uk_conv_nchwc"
+// CHECK-SAME: outs(%[[EMPTY]] :
+// CHECK-SAME: %[[FLAGS]] :
+// CHECK: return %[[MICRO_KERNEL]]#0
+
+// -----
+
+// Dynamic inner tiles (e.g. scalable c0/k0) flow to the ukernel as runtime
+// tensor.dim/index_cast k0/c0 operands instead of baked constants.
+
+#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_nchwc_dynamic_tiles(%input: tensor<1x1x16x16x?xf32>, %filter: tensor<1x1x3x3x?x?xf32>, %output: tensor<1x1x14x14x?xf32>) -> tensor<1x1x14x14x?xf32> attributes {
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz", {ukernels = "all", target_triple = "x86_64-xyz-xyz", cpu_features = "+avx512f"}>
+} {
+ %0 = linalg.generic {indexing_maps = [#map, #map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction", "reduction", "parallel", "reduction"]} ins(%input, %filter : tensor<1x1x16x16x?xf32>, tensor<1x1x3x3x?x?xf32>) outs(%output : tensor<1x1x14x14x?xf32>) {
+ ^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<1x1x14x14x?xf32>
+ return %0 : tensor<1x1x14x14x?xf32>
+}
+// CHECK-LABEL: func @conv_nchwc_dynamic_tiles(
+// CHECK-SAME: %[[INPUT:[a-zA-Z0-9]+]]: tensor<1x1x16x16x?xf32>
+// CHECK-SAME: %[[FILTER:[a-zA-Z0-9]+]]: tensor<1x1x3x3x?x?xf32>
+// CHECK-SAME: %[[OUTPUT:[a-zA-Z0-9]+]]: tensor<1x1x14x14x?xf32>
+// CHECK-DAG: %[[C4:.+]] = arith.constant 4 : index
+// CHECK-DAG: %[[C5:.+]] = arith.constant 5 : index
+// CHECK: %[[K0_DIM:.+]] = tensor.dim %[[FILTER]], %[[C5]]
+// CHECK: %[[K0:.+]] = arith.index_cast %[[K0_DIM]] : index to i32
+// CHECK: %[[C0_DIM:.+]] = tensor.dim %[[FILTER]], %[[C4]]
+// CHECK: %[[C0:.+]] = arith.index_cast %[[C0_DIM]] : index to i32
+// CHECK: iree_codegen.ukernel.generic "iree_uk_conv_nchwc"
+// CHECK-SAME: ins(%[[INPUT]], %[[FILTER]] :
+// CHECK-SAME: outs(%[[OUTPUT]] :
+// CHECK-SAME: %[[K0]], %[[C0]],
+
+// -----
+
+// Without a ukernel target attribute the conv generic is left untouched.
+
+#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 @negative_conv_ukernel(%input: tensor<1x1x16x16x16xf32>, %filter: tensor<1x1x3x3x16x16xf32>, %output: tensor<1x1x14x14x16xf32>) -> tensor<1x1x14x14x16xf32> attributes {
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz", {target_triple = "x86_64-xyz-xyz", cpu_features = "+avx512f"}>
+} {
+ %0 = linalg.generic {indexing_maps = [#map, #map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction", "reduction", "parallel", "reduction"]} ins(%input, %filter : tensor<1x1x16x16x16xf32>, tensor<1x1x3x3x16x16xf32>) outs(%output : 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-LABEL: func @negative_conv_ukernel(
+// CHECK-NOT: iree_uk_conv_nchwc
+// CHECK: linalg.generic
+
+// -----
+
+// linalg.generic with non-MAC body is not recognized as convolution and left as is.
+
+#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 @negative_conv_ukernel_non_conv_body(%input: tensor<1x1x16x16x16xf32>, %filter: tensor<1x1x3x3x16x16xf32>, %output: tensor<1x1x14x14x16xf32>) -> tensor<1x1x14x14x16xf32> attributes {
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz", {ukernels = "all", target_triple = "x86_64-xyz-xyz", cpu_features = "+avx512f"}>
+} {
+ %0 = linalg.generic {indexing_maps = [#map, #map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction", "reduction", "parallel", "reduction"]} ins(%input, %filter : tensor<1x1x16x16x16xf32>, tensor<1x1x3x3x16x16xf32>) outs(%output : tensor<1x1x14x14x16xf32>) {
+ ^bb0(%in: f32, %in_0: f32, %out: f32):
+ %1 = arith.subf %in, %in_0 : f32
+ %2 = arith.addf %1, %out : f32
+ linalg.yield %2 : f32
+ } -> tensor<1x1x14x14x16xf32>
+ return %0 : tensor<1x1x14x14x16xf32>
+}
+// CHECK-LABEL: func @negative_conv_ukernel_non_conv_body(
+// CHECK-NOT: iree_uk_conv_nchwc
+// CHECK: linalg.generic
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/UKernelOps.cpp b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/UKernelOps.cpp
index aec4927..d94a724 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/UKernelOps.cpp
+++ b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/UKernelOps.cpp
@@ -224,7 +224,7 @@
});
auto stridedDimsArrayAttr =
ArrayAttr::get(builder.getContext(), stridedDimsAttrs);
- build(builder, result, outputs.getTypes(), uKernelFnName, inputs, outputs,
+ build(builder, result, resultTypes, uKernelFnName, inputs, outputs,
otherOperands, fnDefAttrs, stridedDimsArrayAttr);
}
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp b/compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp
index 59e5407..ad3f892 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp
@@ -3117,11 +3117,15 @@
getCPUTranslationInfo(op.getContext(), CPUPipeline::Mmt4dTilingExpert));
}
-/// Assigns CPUDoubleTilingExpert pipeline config to the data-tiled
-/// convolution generic.
+/// Assigns lowering config to the data-tiled convolution generic.
static LogicalResult
setConvDataTiledGenericRootConfig(mlir::FunctionOpInterface entryPointFn,
linalg::GenericOp convOp) {
+ IREE::HAL::ExecutableTargetAttr targetAttr =
+ IREE::HAL::ExecutableTargetAttr::lookup(convOp);
+ bool useUkernel =
+ targetAttr && hasUkernel(targetAttr.getConfiguration(), "conv_nchwc");
+
// Loop ranges:
// (d0, d1, d2, d3, d4, d5, d6, d7, d8)
// (n, OC/k0, OH, OW, IC/c0, FH, FW, k0, c0)
@@ -3166,8 +3170,23 @@
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.
+ // When the conv ukernel is enabled we route through `Mmt4dTilingExpert`
+ // rather than a dedicated conv pipeline, because it already runs the modern
+ // lowerings a data-tiled conv dispatch needs: it invokes
+ // `CPULowerToUKernelsPass`, which matches the data-tiled conv
+ // `linalg.generic` and rewrites it to `iree_codegen.ukernel.generic` op,
+ // alongside tile-and-fuse and the generic vectorization passes.
+ // TODO: factor out a data-tiled pipeline that can be shared
+ // between mmt4d/inner_tiled/conv.
+ if (useUkernel) {
+ return setOpConfigAndEntryPointFnTranslation(
+ entryPointFn, convOp, loweringConfig,
+ getCPUTranslationInfo(convOp.getContext(),
+ CPUPipeline::Mmt4dTilingExpert));
+ }
+
+ // Enable loop peeling so the OW tile yields a static-shaped main loop that
+ // vectorizes cleanly, leaving the remainder to a scalar epilogue.
DictionaryAttr pipelineConfig =
getPipelineConfWithPeelingAttr(convOp.getContext());
return setOpConfigAndEntryPointFnTranslation(
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 8fc1c54..919d832 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
@@ -479,6 +479,28 @@
// -----
+#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", ukernels = "conv_nchwc"}>
+#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_ukernel_enable(%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<Mmt4dTilingExpert>>
+// CHECK: func.func @conv_2d_nchwc_data_tiled_ukernel_enable(
+// 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
