[Codegen] Materialize encoding info for convolution with NCHWc layout (#24714)
Implements convolution encoding materialization for data-tiled layouts.
Layouts:
- Input: [N, IC/c0, H, W, c0]
- Filter: [OC/k0, IC/c0, FH, FW, c0, k0]
- Output: [N, OC/k0, OH, OW, k0]
where k0 and c0 are the inner tile sizes for the output and input
channels, respectively.
Scope:
- Only 2D convolutions receive a data-tiled layout.
- Grouped, depthwise, 1D, and 3D convolutions, and non-unit dilations,
all materialize to the identity layout (existing un-tiled path, no
behavior change for these cases).
- Arbitrary strides are supported.
- Supports packed rank-5 input/output and rank-6 filter tensors. Rank-4
input/output is also accepted when `N=1` batch dimension has been folded
by upstream passes.
- Currently supports fp32 convolutions only.
Supported input forms:
- nhwc_hwcf
- nchw_fchw
- nhwc_fhwc
All are materialized to the same packed NCHWc layout above.
Architecture-specific inner tiles:
| Architecture | Condition | Inner tile (OC, IC) | Target ISA |
|--------------|-------------------|---------------------|-----------------|
| x86-64 | float, `+avx512f` | `{16,16}` | `VFMADD*` (zmm) |
| arm64 | float | `{8,8}` | `FMLA` (NEON) |
All other configurations materialize to the identity layout.
Layout motivation can be found here:
https://hackmd.io/@phemashekar/conv-dt-layout
---
Co-authored-by: Jelle Schuhmacher <schuehmacher@roofline.ai>
Co-authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Pooja Hemashekar <hemashekar@roofline.ai>
---------
Signed-off-by: Pooja Hemashekar <hemashekar@roofline.ai>
Co-authored-by: Jelle Schuhmacher <schuehmacher@roofline.ai>
Co-authored-by: Ege Beysel <beyselege@gmail.com>diff --git a/compiler/src/iree/compiler/Codegen/Common/test/BUILD.bazel b/compiler/src/iree/compiler/Codegen/Common/test/BUILD.bazel
index dbb4c96..fbc1e5a 100644
--- a/compiler/src/iree/compiler/Codegen/Common/test/BUILD.bazel
+++ b/compiler/src/iree/compiler/Codegen/Common/test/BUILD.bazel
@@ -88,6 +88,8 @@
"lower_ukernel_to_calls.mlir",
"lowering_config_interpreter.mlir",
"materialize_encoding_aarch64.mlir",
+ "materialize_encoding_conv_aarch64.mlir",
+ "materialize_encoding_conv_x86_64.mlir",
"materialize_encoding_for_iree_ops.mlir",
"materialize_encoding_gfx1100.mlir",
"materialize_encoding_gfx908.mlir",
diff --git a/compiler/src/iree/compiler/Codegen/Common/test/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/Common/test/CMakeLists.txt
index c950ca2..f3220e2 100644
--- a/compiler/src/iree/compiler/Codegen/Common/test/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Codegen/Common/test/CMakeLists.txt
@@ -83,6 +83,8 @@
"lower_ukernel_to_calls.mlir"
"lowering_config_interpreter.mlir"
"materialize_encoding_aarch64.mlir"
+ "materialize_encoding_conv_aarch64.mlir"
+ "materialize_encoding_conv_x86_64.mlir"
"materialize_encoding_for_iree_ops.mlir"
"materialize_encoding_gfx1100.mlir"
"materialize_encoding_gfx908.mlir"
diff --git a/compiler/src/iree/compiler/Codegen/Common/test/materialize_encoding_conv_aarch64.mlir b/compiler/src/iree/compiler/Codegen/Common/test/materialize_encoding_conv_aarch64.mlir
new file mode 100644
index 0000000..48e00bc
--- /dev/null
+++ b/compiler/src/iree/compiler/Codegen/Common/test/materialize_encoding_conv_aarch64.mlir
@@ -0,0 +1,574 @@
+// RUN: iree-opt --pass-pipeline="builtin.module(func.func(iree-codegen-materialize-device-encoding))" --split-input-file %s | FileCheck %s
+
+// Map invariant materialization (all input formats must produce the same canonical generics):
+// input: (d0..d8) -> (d0, d4, d2 + d5, d3 + d6, d8)
+// filter: (d0..d8) -> (d1, d4, d5, d6, d8, d7)
+// output: (d0..d8) -> (d0, d1, d2, d3, d7)
+
+// conv_2d_nhwc_hwcf materialization.
+
+// Pack input from [N, H, W, C] -> [N, IC/c0, H, W, c0].
+
+#nhwc_hwcf_map_in = affine_map<(n, oh, ow, oc, fh, fw, ic) -> (n, oh + fh, ow + fw, ic)>
+#nhwc_hwcf_map_f = affine_map<(n, oh, ow, oc, fh, fw, ic) -> (fh, fw, ic, oc)>
+#nhwc_hwcf_map_out = affine_map<(n, oh, ow, oc, fh, fw, ic) -> (n, oh, ow, oc)>
+
+// CHECK: #[[$M_IN:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d0, d4, d2 + d5, d3 + d6, d8)>
+// CHECK: #[[$M_FLT:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d1, d4, d5, d6, d8, d7)>
+// CHECK: #[[$M_OUT:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d0, d1, d2, d3, d7)>
+
+#encoding_nhwc_hwcf_input = #iree_encoding.encoding<operand_index = 0, op_type = conv,
+ element_types = [f32, f32, f32],
+ user_indexing_maps = [#nhwc_hwcf_map_in, #nhwc_hwcf_map_f, #nhwc_hwcf_map_out],
+ iteration_sizes = [1, 14, 14, 8, 3, 3, 4]>
+
+func.func @conv_input_pack(%arg0: tensor<1x16x16x4xf32>)
+ -> tensor<1x16x16x4xf32, #encoding_nhwc_hwcf_input>
+ attributes {
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz",
+ {target_triple="aarch64-xyz-xyz", cpu_features="+neon",
+ iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+} {
+ %0 = iree_encoding.set_encoding %arg0
+ : tensor<1x16x16x4xf32> -> tensor<1x16x16x4xf32, #encoding_nhwc_hwcf_input>
+ return %0 : tensor<1x16x16x4xf32, #encoding_nhwc_hwcf_input>
+}
+// CHECK-LABEL: func.func @conv_input_pack
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<1x16x16x4xf32>
+// CHECK: %[[PACK:.+]] = linalg.pack %[[ARG0]]
+// CHECK-SAME: outer_dims_perm = [0, 3, 1, 2]
+// CHECK-SAME: inner_dims_pos = [3]
+// CHECK-SAME: inner_tiles = [8]
+// CHECK-SAME: : tensor<1x16x16x4xf32> -> tensor<1x1x16x16x8xf32>
+// CHECK: return %[[PACK]]
+
+
+// Pack filter from [KH, KW, C, F] -> [OC/k0, IC/c0, KH, KW, c0, k0].
+
+#encoding_nhwc_hwcf_filter = #iree_encoding.encoding<operand_index = 1, op_type = conv,
+ element_types = [f32, f32, f32],
+ user_indexing_maps = [#nhwc_hwcf_map_in, #nhwc_hwcf_map_f, #nhwc_hwcf_map_out],
+ iteration_sizes = [1, 14, 14, 8, 3, 3, 4]>
+
+func.func @conv_filter_pack(%arg0: tensor<3x3x4x8xf32>)
+ -> tensor<3x3x4x8xf32, #encoding_nhwc_hwcf_filter>
+ attributes {
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz",
+ {target_triple="aarch64-xyz-xyz", cpu_features="+neon",
+ iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+} {
+ %0 = iree_encoding.set_encoding %arg0
+ : tensor<3x3x4x8xf32> -> tensor<3x3x4x8xf32, #encoding_nhwc_hwcf_filter>
+ return %0 : tensor<3x3x4x8xf32, #encoding_nhwc_hwcf_filter>
+}
+// CHECK-LABEL: func.func @conv_filter_pack
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<3x3x4x8xf32>
+// CHECK: %[[PACK:.+]] = linalg.pack %[[ARG0]]
+// CHECK-SAME: outer_dims_perm = [3, 2, 0, 1]
+// CHECK-SAME: inner_dims_pos = [2, 3]
+// CHECK-SAME: inner_tiles = [8, 8]
+// CHECK-SAME: : tensor<3x3x4x8xf32> -> tensor<1x1x3x3x8x8xf32>
+// CHECK: return %[[PACK]]
+
+
+// Unpack output from [N, OC/k0, OH, OW, k0] -> [N, OH, OW, OC].
+
+#encoding_nhwc_hwcf_output = #iree_encoding.encoding<operand_index = 2, op_type = conv,
+ element_types = [f32, f32, f32],
+ user_indexing_maps = [#nhwc_hwcf_map_in, #nhwc_hwcf_map_f, #nhwc_hwcf_map_out],
+ iteration_sizes = [1, 14, 14, 8, 3, 3, 4]>
+
+func.func @conv_output_unset(%arg0: tensor<1x14x14x8xf32, #encoding_nhwc_hwcf_output>)
+ -> tensor<1x14x14x8xf32>
+ attributes {
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz",
+ {target_triple="aarch64-xyz-xyz", cpu_features="+neon",
+ iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+} {
+ %0 = iree_encoding.unset_encoding %arg0
+ : tensor<1x14x14x8xf32, #encoding_nhwc_hwcf_output> -> tensor<1x14x14x8xf32>
+ return %0 : tensor<1x14x14x8xf32>
+}
+// CHECK-LABEL: func.func @conv_output_unset
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<1x1x14x14x8xf32>
+// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[ARG0]]
+// CHECK-SAME: outer_dims_perm = [0, 3, 1, 2]
+// CHECK-SAME: inner_dims_pos = [3]
+// CHECK-SAME: inner_tiles = [8]
+// CHECK-SAME: : tensor<1x1x14x14x8xf32> -> tensor<1x14x14x8xf32>
+// CHECK: return %[[UNPACK]]
+
+
+// Conv materialization.
+
+func.func @conv2d_nhwc_hwcf_materialize(
+ %input : tensor<1x16x16x4xf32, #encoding_nhwc_hwcf_input>,
+ %filter : tensor<3x3x4x8xf32, #encoding_nhwc_hwcf_filter>,
+ %output : tensor<1x14x14x8xf32, #encoding_nhwc_hwcf_output>)
+ -> tensor<1x14x14x8xf32, #encoding_nhwc_hwcf_output>
+ attributes {
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz",
+ {target_triple="aarch64-xyz-xyz", cpu_features="+neon",
+ iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+} {
+ %0 = linalg.conv_2d_nhwc_hwcf
+ {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
+ ins(%input, %filter
+ : tensor<1x16x16x4xf32, #encoding_nhwc_hwcf_input>,
+ tensor<3x3x4x8xf32, #encoding_nhwc_hwcf_filter>)
+ outs(%output : tensor<1x14x14x8xf32, #encoding_nhwc_hwcf_output>)
+ -> tensor<1x14x14x8xf32, #encoding_nhwc_hwcf_output>
+ return %0 : tensor<1x14x14x8xf32, #encoding_nhwc_hwcf_output>
+}
+// CHECK-LABEL: func.func @conv2d_nhwc_hwcf_materialize
+// CHECK-SAME: %[[INPUT:.+]]: tensor<1x1x16x16x8xf32>
+// CHECK-SAME: %[[FILTER:.+]]: tensor<1x1x3x3x8x8xf32>
+// CHECK-SAME: %[[OUTPUT:.+]]: tensor<1x1x14x14x8xf32>
+//
+// 9D data-tiled conv generic with three operands (input, filter, output):
+// CHECK: %[[RESULT:.+]] = linalg.generic
+// CHECK-SAME: indexing_maps = [#[[$M_IN]], #[[$M_FLT]], #[[$M_OUT]]]
+// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel",
+// CHECK-SAME: "reduction", "reduction", "reduction",
+// CHECK-SAME: "parallel", "reduction"]
+// CHECK-SAME: ins(%[[INPUT]], %[[FILTER]]
+// CHECK-SAME: outs(%[[OUTPUT]]
+// CHECK: %[[MUL:.+]] = arith.mulf
+// CHECK: %[[ADD:.+]] = arith.addf %[[MUL]]
+// CHECK: linalg.yield %[[ADD]]
+// CHECK: return %[[RESULT]]
+
+
+// Full conv materialization.
+
+func.func @conv2d_nhwc_hwcf_materialize_full(
+ %input : tensor<1x16x16x4xf32>,
+ %filter : tensor<3x3x4x8xf32>,
+ %output : tensor<1x14x14x8xf32>)
+ -> tensor<1x14x14x8xf32>
+ attributes {
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz",
+ {target_triple="aarch64-xyz-xyz", cpu_features="+neon",
+ iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+} {
+ %enc_in = iree_encoding.set_encoding %input
+ : tensor<1x16x16x4xf32> -> tensor<1x16x16x4xf32, #encoding_nhwc_hwcf_input>
+ %enc_f = iree_encoding.set_encoding %filter
+ : tensor<3x3x4x8xf32> -> tensor<3x3x4x8xf32, #encoding_nhwc_hwcf_filter>
+ %enc_out = iree_encoding.set_encoding %output
+ : tensor<1x14x14x8xf32> -> tensor<1x14x14x8xf32, #encoding_nhwc_hwcf_output>
+ %0 = linalg.conv_2d_nhwc_hwcf
+ {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
+ ins(%enc_in, %enc_f
+ : tensor<1x16x16x4xf32, #encoding_nhwc_hwcf_input>,
+ tensor<3x3x4x8xf32, #encoding_nhwc_hwcf_filter>)
+ outs(%enc_out : tensor<1x14x14x8xf32, #encoding_nhwc_hwcf_output>)
+ -> tensor<1x14x14x8xf32, #encoding_nhwc_hwcf_output>
+ %1 = iree_encoding.unset_encoding %0
+ : tensor<1x14x14x8xf32, #encoding_nhwc_hwcf_output> -> tensor<1x14x14x8xf32>
+ return %1 : tensor<1x14x14x8xf32>
+}
+// CHECK-LABEL: func.func @conv2d_nhwc_hwcf_materialize_full
+// CHECK-SAME: %[[IN:[a-zA-Z0-9]+]]: tensor<1x16x16x4xf32>
+// CHECK-SAME: %[[F:[a-zA-Z0-9]+]]: tensor<3x3x4x8xf32>
+// CHECK-SAME: %[[OUT:[a-zA-Z0-9]+]]: tensor<1x14x14x8xf32>
+//
+// CHECK: %[[PACK_IN:.+]] = linalg.pack %[[IN]]
+// CHECK-SAME: outer_dims_perm = [0, 3, 1, 2]
+// CHECK-SAME: inner_dims_pos = [3]
+// CHECK-SAME: inner_tiles = [8]
+// CHECK-SAME: : tensor<1x16x16x4xf32> -> tensor<1x1x16x16x8xf32>
+//
+// CHECK: %[[PACK_F:.+]] = linalg.pack %[[F]]
+// CHECK-SAME: outer_dims_perm = [3, 2, 0, 1]
+// CHECK-SAME: inner_dims_pos = [2, 3]
+// CHECK-SAME: inner_tiles = [8, 8]
+// CHECK-SAME: : tensor<3x3x4x8xf32> -> tensor<1x1x3x3x8x8xf32>
+//
+// CHECK: %[[PACK_OUT:.+]] = linalg.pack %[[OUT]]
+// CHECK-SAME: outer_dims_perm = [0, 3, 1, 2]
+// CHECK-SAME: inner_dims_pos = [3]
+// CHECK-SAME: inner_tiles = [8]
+// CHECK-SAME: : tensor<1x14x14x8xf32> -> tensor<1x1x14x14x8xf32>
+//
+// CHECK: %[[RESULT:.+]] = linalg.generic
+// CHECK-SAME: ins(%[[PACK_IN]], %[[PACK_F]]
+// CHECK-SAME: outs(%[[PACK_OUT]]
+//
+// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[RESULT]]
+// CHECK-SAME: outer_dims_perm = [0, 3, 1, 2]
+// CHECK-SAME: inner_dims_pos = [3]
+// CHECK-SAME: inner_tiles = [8]
+// CHECK-SAME: : tensor<1x1x14x14x8xf32> -> tensor<1x14x14x8xf32>
+
+
+// NHWC-HWCF conv expressed as a linalg.generic carrying conv encodings.
+// It must lower to the same pack / 9D generic / unpack as the named conv2d_nhwc_hwcf_materialize_full above.
+
+func.func @conv2d_nhwc_hwcf_materialize_generic(
+ %input : tensor<1x16x16x4xf32>,
+ %filter : tensor<3x3x4x8xf32>,
+ %output : tensor<1x14x14x8xf32>)
+ -> tensor<1x14x14x8xf32>
+ attributes {
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz",
+ {target_triple="aarch64-xyz-xyz", cpu_features="+neon",
+ iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+} {
+ %enc_in = iree_encoding.set_encoding %input
+ : tensor<1x16x16x4xf32> -> tensor<1x16x16x4xf32, #encoding_nhwc_hwcf_input>
+ %enc_f = iree_encoding.set_encoding %filter
+ : tensor<3x3x4x8xf32> -> tensor<3x3x4x8xf32, #encoding_nhwc_hwcf_filter>
+ %enc_out = iree_encoding.set_encoding %output
+ : tensor<1x14x14x8xf32> -> tensor<1x14x14x8xf32, #encoding_nhwc_hwcf_output>
+ %0 = linalg.generic
+ {indexing_maps = [#nhwc_hwcf_map_in, #nhwc_hwcf_map_f, #nhwc_hwcf_map_out],
+ iterator_types = ["parallel", "parallel", "parallel", "parallel",
+ "reduction", "reduction", "reduction"]}
+ ins(%enc_in, %enc_f
+ : tensor<1x16x16x4xf32, #encoding_nhwc_hwcf_input>,
+ tensor<3x3x4x8xf32, #encoding_nhwc_hwcf_filter>)
+ outs(%enc_out : tensor<1x14x14x8xf32, #encoding_nhwc_hwcf_output>) {
+ ^bb0(%in: f32, %f: f32, %acc: f32):
+ %m = arith.mulf %in, %f : f32
+ %a = arith.addf %acc, %m : f32
+ linalg.yield %a : f32
+ } -> tensor<1x14x14x8xf32, #encoding_nhwc_hwcf_output>
+ %1 = iree_encoding.unset_encoding %0
+ : tensor<1x14x14x8xf32, #encoding_nhwc_hwcf_output> -> tensor<1x14x14x8xf32>
+ return %1 : tensor<1x14x14x8xf32>
+}
+// CHECK-LABEL: func.func @conv2d_nhwc_hwcf_materialize_generic
+// CHECK-SAME: %[[IN:[a-zA-Z0-9]+]]: tensor<1x16x16x4xf32>
+// CHECK-SAME: %[[F:[a-zA-Z0-9]+]]: tensor<3x3x4x8xf32>
+// CHECK-SAME: %[[OUT:[a-zA-Z0-9]+]]: tensor<1x14x14x8xf32>
+//
+// CHECK: %[[PACK_IN:.+]] = linalg.pack %[[IN]]
+// CHECK-SAME: outer_dims_perm = [0, 3, 1, 2]
+// CHECK-SAME: inner_dims_pos = [3]
+// CHECK-SAME: inner_tiles = [8]
+// CHECK-SAME: : tensor<1x16x16x4xf32> -> tensor<1x1x16x16x8xf32>
+//
+// CHECK: %[[PACK_F:.+]] = linalg.pack %[[F]]
+// CHECK-SAME: outer_dims_perm = [3, 2, 0, 1]
+// CHECK-SAME: inner_dims_pos = [2, 3]
+// CHECK-SAME: inner_tiles = [8, 8]
+// CHECK-SAME: : tensor<3x3x4x8xf32> -> tensor<1x1x3x3x8x8xf32>
+//
+// CHECK: %[[PACK_OUT:.+]] = linalg.pack %[[OUT]]
+// CHECK-SAME: outer_dims_perm = [0, 3, 1, 2]
+// CHECK-SAME: inner_dims_pos = [3]
+// CHECK-SAME: inner_tiles = [8]
+// CHECK-SAME: : tensor<1x14x14x8xf32> -> tensor<1x1x14x14x8xf32>
+//
+// CHECK: %[[RESULT:.+]] = linalg.generic
+// CHECK-SAME: ins(%[[PACK_IN]], %[[PACK_F]]
+// CHECK-SAME: outs(%[[PACK_OUT]]
+//
+// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[RESULT]]
+// CHECK-SAME: outer_dims_perm = [0, 3, 1, 2]
+// CHECK-SAME: inner_dims_pos = [3]
+// CHECK-SAME: inner_tiles = [8]
+// CHECK-SAME: : tensor<1x1x14x14x8xf32> -> tensor<1x14x14x8xf32>
+
+
+// -----
+
+// conv_2d_nchw_fchw materialization.
+
+// Pack input from [N, C, H, W] -> [N, IC/c0, H, W, c0].
+
+#nchw_fchw_map_in = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d4, d2 + d5, d3 + d6)>
+#nchw_fchw_map_f = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d1, d4, d5, d6)>
+#nchw_fchw_map_out = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3)>
+
+// CHECK: #[[$M_IN:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d0, d4, d2 + d5, d3 + d6, d8)>
+// CHECK: #[[$M_FLT:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d1, d4, d5, d6, d8, d7)>
+// CHECK: #[[$M_OUT:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d0, d1, d2, d3, d7)>
+
+#encoding_nhwc_fchw_input = #iree_encoding.encoding<operand_index = 0, op_type = conv,
+ element_types = [f32, f32, f32],
+ user_indexing_maps = [#nchw_fchw_map_in, #nchw_fchw_map_f, #nchw_fchw_map_out],
+ iteration_sizes = [1, 8, 14, 14, 4, 3, 3]>
+
+func.func @conv_nchw_input_pack(%arg0: tensor<1x4x16x16xf32>)
+ -> tensor<1x4x16x16xf32, #encoding_nhwc_fchw_input>
+ attributes {
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz",
+ {target_triple="aarch64-xyz-xyz", cpu_features="+neon",
+ iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+} {
+ %0 = iree_encoding.set_encoding %arg0
+ : tensor<1x4x16x16xf32> -> tensor<1x4x16x16xf32, #encoding_nhwc_fchw_input>
+ return %0 : tensor<1x4x16x16xf32, #encoding_nhwc_fchw_input>
+}
+// CHECK-LABEL: func.func @conv_nchw_input_pack
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<1x4x16x16xf32>
+// CHECK: %[[PACK:.+]] = linalg.pack %[[ARG0]]
+// CHECK-SAME: outer_dims_perm = [0, 1, 2, 3]
+// CHECK-SAME: inner_dims_pos = [1]
+// CHECK-SAME: inner_tiles = [8]
+// CHECK-SAME: : tensor<1x4x16x16xf32> -> tensor<1x1x16x16x8xf32>
+// CHECK: return %[[PACK]]
+
+
+// Pack filter from [F, C, KH, KW] -> [OC/k0, IC/c0, KH, KW, c0, k0].
+
+#encoding_nhwc_fchw_filter = #iree_encoding.encoding<operand_index = 1, op_type = conv,
+ element_types = [f32, f32, f32],
+ user_indexing_maps = [#nchw_fchw_map_in, #nchw_fchw_map_f, #nchw_fchw_map_out],
+ iteration_sizes = [1, 8, 14, 14, 4, 3, 3]>
+
+func.func @conv_fchw_filter_pack(%arg0: tensor<8x4x3x3xf32>)
+ -> tensor<8x4x3x3xf32, #encoding_nhwc_fchw_filter>
+ attributes {
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz",
+ {target_triple="aarch64-xyz-xyz", cpu_features="+neon",
+ iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+} {
+ %0 = iree_encoding.set_encoding %arg0
+ : tensor<8x4x3x3xf32> -> tensor<8x4x3x3xf32, #encoding_nhwc_fchw_filter>
+ return %0 : tensor<8x4x3x3xf32, #encoding_nhwc_fchw_filter>
+}
+// CHECK-LABEL: func.func @conv_fchw_filter_pack
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<8x4x3x3xf32>
+// CHECK: %[[PACK:.+]] = linalg.pack %[[ARG0]]
+// CHECK-SAME: outer_dims_perm = [0, 1, 2, 3]
+// CHECK-SAME: inner_dims_pos = [1, 0]
+// CHECK-SAME: inner_tiles = [8, 8]
+// CHECK-SAME: : tensor<8x4x3x3xf32> -> tensor<1x1x3x3x8x8xf32>
+// CHECK: return %[[PACK]]
+
+
+// Unpack output from [N, OC/k0, OH, OW, k0] -> [N, OC, OH, OW].
+
+#encoding_nhwc_fchw_output = #iree_encoding.encoding<operand_index = 2, op_type = conv,
+ element_types = [f32, f32, f32],
+ user_indexing_maps = [#nchw_fchw_map_in, #nchw_fchw_map_f, #nchw_fchw_map_out],
+ iteration_sizes = [1, 8, 14, 14, 4, 3, 3]>
+
+func.func @conv_nchw_output_unset(%arg0: tensor<1x8x14x14xf32, #encoding_nhwc_fchw_output>)
+ -> tensor<1x8x14x14xf32>
+ attributes {
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz",
+ {target_triple="aarch64-xyz-xyz", cpu_features="+neon",
+ iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+} {
+ %0 = iree_encoding.unset_encoding %arg0
+ : tensor<1x8x14x14xf32, #encoding_nhwc_fchw_output> -> tensor<1x8x14x14xf32>
+ return %0 : tensor<1x8x14x14xf32>
+}
+// CHECK-LABEL: func.func @conv_nchw_output_unset
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<1x1x14x14x8xf32>
+// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[ARG0]]
+// CHECK-SAME: outer_dims_perm = [0, 1, 2, 3]
+// CHECK-SAME: inner_dims_pos = [1]
+// CHECK-SAME: inner_tiles = [8]
+// CHECK-SAME: : tensor<1x1x14x14x8xf32> -> tensor<1x8x14x14xf32>
+// CHECK: return %[[UNPACK]]
+
+
+// Conv materialization.
+
+func.func @conv2d_nchw_fchw_materialize(
+ %input : tensor<1x4x16x16xf32, #encoding_nhwc_fchw_input>,
+ %filter : tensor<8x4x3x3xf32, #encoding_nhwc_fchw_filter>,
+ %output : tensor<1x8x14x14xf32, #encoding_nhwc_fchw_output>)
+ -> tensor<1x8x14x14xf32, #encoding_nhwc_fchw_output>
+ attributes {
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz",
+ {target_triple="aarch64-xyz-xyz", cpu_features="+neon",
+ iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+} {
+ %0 = linalg.conv_2d_nchw_fchw
+ {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
+ ins(%input, %filter
+ : tensor<1x4x16x16xf32, #encoding_nhwc_fchw_input>,
+ tensor<8x4x3x3xf32, #encoding_nhwc_fchw_filter>)
+ outs(%output : tensor<1x8x14x14xf32, #encoding_nhwc_fchw_output>)
+ -> tensor<1x8x14x14xf32, #encoding_nhwc_fchw_output>
+ return %0 : tensor<1x8x14x14xf32, #encoding_nhwc_fchw_output>
+}
+// CHECK-LABEL: func.func @conv2d_nchw_fchw_materialize
+// CHECK-SAME: %[[INPUT:.+]]: tensor<1x1x16x16x8xf32>
+// CHECK-SAME: %[[FILTER:.+]]: tensor<1x1x3x3x8x8xf32>
+// CHECK-SAME: %[[OUTPUT:.+]]: tensor<1x1x14x14x8xf32>
+//
+// CHECK: %[[RESULT:.+]] = linalg.generic
+// CHECK-SAME: indexing_maps = [#[[$M_IN]], #[[$M_FLT]], #[[$M_OUT]]]
+// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel",
+// CHECK-SAME: "reduction", "reduction", "reduction",
+// CHECK-SAME: "parallel", "reduction"]
+// CHECK-SAME: ins(%[[INPUT]], %[[FILTER]]
+// CHECK-SAME: outs(%[[OUTPUT]]
+// CHECK: %[[MUL:.+]] = arith.mulf
+// CHECK: %[[ADD:.+]] = arith.addf %[[MUL]]
+// CHECK: linalg.yield %[[ADD]]
+// CHECK: return %[[RESULT]]
+
+
+// Full conv materialization.
+
+func.func @conv2d_nchw_fchw_materialize_full(
+ %input : tensor<1x4x16x16xf32>,
+ %filter : tensor<8x4x3x3xf32>,
+ %output : tensor<1x8x14x14xf32>)
+ -> tensor<1x8x14x14xf32>
+ attributes {
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz",
+ {target_triple="aarch64-xyz-xyz", cpu_features="+neon",
+ iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+} {
+ %enc_in = iree_encoding.set_encoding %input
+ : tensor<1x4x16x16xf32> -> tensor<1x4x16x16xf32, #encoding_nhwc_fchw_input>
+ %enc_f = iree_encoding.set_encoding %filter
+ : tensor<8x4x3x3xf32> -> tensor<8x4x3x3xf32, #encoding_nhwc_fchw_filter>
+ %enc_out = iree_encoding.set_encoding %output
+ : tensor<1x8x14x14xf32> -> tensor<1x8x14x14xf32, #encoding_nhwc_fchw_output>
+ %0 = linalg.conv_2d_nchw_fchw
+ {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
+ ins(%enc_in, %enc_f
+ : tensor<1x4x16x16xf32, #encoding_nhwc_fchw_input>,
+ tensor<8x4x3x3xf32, #encoding_nhwc_fchw_filter>)
+ outs(%enc_out : tensor<1x8x14x14xf32, #encoding_nhwc_fchw_output>)
+ -> tensor<1x8x14x14xf32, #encoding_nhwc_fchw_output>
+ %1 = iree_encoding.unset_encoding %0
+ : tensor<1x8x14x14xf32, #encoding_nhwc_fchw_output> -> tensor<1x8x14x14xf32>
+ return %1 : tensor<1x8x14x14xf32>
+}
+// CHECK-LABEL: func.func @conv2d_nchw_fchw_materialize_full
+// CHECK-SAME: %[[IN:[a-zA-Z0-9]+]]: tensor<1x4x16x16xf32>
+// CHECK-SAME: %[[F:[a-zA-Z0-9]+]]: tensor<8x4x3x3xf32>
+// CHECK-SAME: %[[OUT:[a-zA-Z0-9]+]]: tensor<1x8x14x14xf32>
+//
+// CHECK: %[[PACK_IN:.+]] = linalg.pack %[[IN]]
+// CHECK-SAME: outer_dims_perm = [0, 1, 2, 3]
+// CHECK-SAME: inner_dims_pos = [1]
+// CHECK-SAME: inner_tiles = [8]
+// CHECK-SAME: : tensor<1x4x16x16xf32> -> tensor<1x1x16x16x8xf32>
+//
+// CHECK: %[[PACK_F:.+]] = linalg.pack %[[F]]
+// CHECK-SAME: outer_dims_perm = [0, 1, 2, 3]
+// CHECK-SAME: inner_dims_pos = [1, 0]
+// CHECK-SAME: inner_tiles = [8, 8]
+// CHECK-SAME: : tensor<8x4x3x3xf32> -> tensor<1x1x3x3x8x8xf32>
+//
+// CHECK: %[[PACK_OUT:.+]] = linalg.pack %[[OUT]]
+// CHECK-SAME: outer_dims_perm = [0, 1, 2, 3]
+// CHECK-SAME: inner_dims_pos = [1]
+// CHECK-SAME: inner_tiles = [8]
+// CHECK-SAME: : tensor<1x8x14x14xf32> -> tensor<1x1x14x14x8xf32>
+//
+// CHECK: %[[RESULT:.+]] = linalg.generic
+// CHECK-SAME: ins(%[[PACK_IN]], %[[PACK_F]]
+// CHECK-SAME: outs(%[[PACK_OUT]]
+//
+// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[RESULT]]
+// CHECK-SAME: outer_dims_perm = [0, 1, 2, 3]
+// CHECK-SAME: inner_dims_pos = [1]
+// CHECK-SAME: inner_tiles = [8]
+// CHECK-SAME: : tensor<1x1x14x14x8xf32> -> tensor<1x8x14x14xf32>
+
+// -----
+
+// conv_2d_nhwc_fhwc materialization.
+
+// Pack filter from [F, KH, KW, C] -> [OC/k0, IC/c0, KH, KW, c0, k0].
+
+#nhwc_fhwc_map_in = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1 + d4, d2 + d5, d6)>
+#nhwc_fhwc_map_f = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d3, d4, d5, d6)>
+#nhwc_fhwc_map_out = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3)>
+
+#encoding_nhwc_fhwc_filter = #iree_encoding.encoding<operand_index = 1, op_type = conv,
+ element_types = [f32, f32, f32],
+ user_indexing_maps = [#nhwc_fhwc_map_in, #nhwc_fhwc_map_f, #nhwc_fhwc_map_out],
+ iteration_sizes = [1, 14, 14, 8, 3, 3, 4]>
+
+func.func @conv_fhwc_filter_pack(%arg0: tensor<8x3x3x4xf32>)
+ -> tensor<8x3x3x4xf32, #encoding_nhwc_fhwc_filter>
+ attributes {
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz",
+ {target_triple="aarch64-xyz-xyz", cpu_features="+neon",
+ iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+} {
+ %0 = iree_encoding.set_encoding %arg0
+ : tensor<8x3x3x4xf32> -> tensor<8x3x3x4xf32, #encoding_nhwc_fhwc_filter>
+ return %0 : tensor<8x3x3x4xf32, #encoding_nhwc_fhwc_filter>
+}
+// CHECK-LABEL: func.func @conv_fhwc_filter_pack
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<8x3x3x4xf32>
+// CHECK: %[[PACK:.+]] = linalg.pack %[[ARG0]]
+// CHECK-SAME: outer_dims_perm = [0, 3, 1, 2]
+// CHECK-SAME: inner_dims_pos = [3, 0]
+// CHECK-SAME: inner_tiles = [8, 8]
+// CHECK-SAME: : tensor<8x3x3x4xf32> -> tensor<1x1x3x3x8x8xf32>
+// CHECK: return %[[PACK]]
+
+// -----
+
+// Unit batch-dim folded conv: N=1 stripped upstream, so operands pack to rank-4
+// and the lowering restores the batch via expand/collapse around the 9-D generic.
+
+#bl_map_in = affine_map<(oh, ow, oc, fh, fw, ic) -> (oh + fh, ow + fw, ic)>
+#bl_map_f = affine_map<(oh, ow, oc, fh, fw, ic) -> (fh, fw, ic, oc)>
+#bl_map_out = affine_map<(oh, ow, oc, fh, fw, ic) -> (oh, ow, oc)>
+
+#encoding_bl_input = #iree_encoding.encoding<operand_index = 0, op_type = conv,
+ element_types = [f32, f32, f32],
+ user_indexing_maps = [#bl_map_in, #bl_map_f, #bl_map_out],
+ iteration_sizes = [14, 14, 8, 3, 3, 4]>
+#encoding_bl_filter = #iree_encoding.encoding<operand_index = 1, op_type = conv,
+ element_types = [f32, f32, f32],
+ user_indexing_maps = [#bl_map_in, #bl_map_f, #bl_map_out],
+ iteration_sizes = [14, 14, 8, 3, 3, 4]>
+#encoding_bl_output = #iree_encoding.encoding<operand_index = 2, op_type = conv,
+ element_types = [f32, f32, f32],
+ user_indexing_maps = [#bl_map_in, #bl_map_f, #bl_map_out],
+ iteration_sizes = [14, 14, 8, 3, 3, 4]>
+
+func.func @conv2d_batchless_materialize(
+ %input : tensor<16x16x4xf32, #encoding_bl_input>,
+ %filter : tensor<3x3x4x8xf32, #encoding_bl_filter>,
+ %output : tensor<14x14x8xf32, #encoding_bl_output>)
+ -> tensor<14x14x8xf32, #encoding_bl_output>
+ attributes {
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz",
+ {target_triple="aarch64-xyz-xyz", cpu_features="+neon",
+ iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+} {
+ %0 = linalg.generic
+ {indexing_maps = [#bl_map_in, #bl_map_f, #bl_map_out],
+ iterator_types = ["parallel", "parallel", "parallel",
+ "reduction", "reduction", "reduction"]}
+ ins(%input, %filter
+ : tensor<16x16x4xf32, #encoding_bl_input>,
+ tensor<3x3x4x8xf32, #encoding_bl_filter>)
+ outs(%output : tensor<14x14x8xf32, #encoding_bl_output>) {
+ ^bb0(%in: f32, %f: f32, %acc: f32):
+ %m = arith.mulf %in, %f : f32
+ %a = arith.addf %acc, %m : f32
+ linalg.yield %a : f32
+ } -> tensor<14x14x8xf32, #encoding_bl_output>
+ return %0 : tensor<14x14x8xf32, #encoding_bl_output>
+}
+// CHECK-LABEL: func.func @conv2d_batchless_materialize
+// CHECK-SAME: %[[INPUT:.+]]: tensor<1x16x16x8xf32>
+// CHECK-SAME: %[[FILTER:.+]]: tensor<1x1x3x3x8x8xf32>
+// CHECK-SAME: %[[OUTPUT:.+]]: tensor<1x14x14x8xf32>
+// CHECK: %[[EXP_IN:.+]] = tensor.expand_shape %[[INPUT]] {{\[}}[0, 1], [2], [3], [4]{{\]}}
+// CHECK-SAME: : tensor<1x16x16x8xf32> into tensor<1x1x16x16x8xf32>
+// CHECK: %[[EXP_OUT:.+]] = tensor.expand_shape %[[OUTPUT]] {{\[}}[0, 1], [2], [3], [4]{{\]}}
+// CHECK-SAME: : tensor<1x14x14x8xf32> into tensor<1x1x14x14x8xf32>
+// CHECK: %[[RESULT:.+]] = linalg.generic
+// CHECK-SAME: indexing_maps = [#[[$M_IN]], #[[$M_FLT]], #[[$M_OUT]]]
+// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel",
+// CHECK-SAME: "reduction", "reduction", "reduction",
+// CHECK-SAME: "parallel", "reduction"]
+// CHECK-SAME: ins(%[[EXP_IN]], %[[FILTER]]
+// CHECK-SAME: outs(%[[EXP_OUT]]
+// CHECK: %[[MUL:.+]] = arith.mulf
+// CHECK: %[[ADD:.+]] = arith.addf %[[MUL]]
+// CHECK: linalg.yield %[[ADD]]
+// CHECK: %[[COLLAPSE:.+]] = tensor.collapse_shape %[[RESULT]] {{\[}}[0, 1], [2], [3], [4]{{\]}}
+// CHECK-SAME: : tensor<1x1x14x14x8xf32> into tensor<1x14x14x8xf32>
+// CHECK: return %[[COLLAPSE]]
diff --git a/compiler/src/iree/compiler/Codegen/Common/test/materialize_encoding_conv_x86_64.mlir b/compiler/src/iree/compiler/Codegen/Common/test/materialize_encoding_conv_x86_64.mlir
new file mode 100644
index 0000000..78101cc
--- /dev/null
+++ b/compiler/src/iree/compiler/Codegen/Common/test/materialize_encoding_conv_x86_64.mlir
@@ -0,0 +1,806 @@
+// RUN: iree-opt --pass-pipeline="builtin.module(func.func(iree-codegen-materialize-device-encoding))" --split-input-file %s | FileCheck %s
+
+// Map invariant materialization (all input formats must produce the same canonical generics):
+// input: (d0..d8) -> (d0, d4, d2 + d5, d3 + d6, d8)
+// filter: (d0..d8) -> (d1, d4, d5, d6, d8, d7)
+// output: (d0..d8) -> (d0, d1, d2, d3, d7)
+
+// conv_2d_nhwc_hwcf materialization.
+
+// Pack input from [N, H, W, C] -> [N, IC/c0, H, W, c0].
+
+#nhwc_hwcf_map_in = affine_map<(n, oh, ow, oc, fh, fw, ic) -> (n, oh + fh, ow + fw, ic)>
+#nhwc_hwcf_map_f = affine_map<(n, oh, ow, oc, fh, fw, ic) -> (fh, fw, ic, oc)>
+#nhwc_hwcf_map_out = affine_map<(n, oh, ow, oc, fh, fw, ic) -> (n, oh, ow, oc)>
+
+// CHECK: #[[$M_IN:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d0, d4, d2 + d5, d3 + d6, d8)>
+// CHECK: #[[$M_FLT:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d1, d4, d5, d6, d8, d7)>
+// CHECK: #[[$M_OUT:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d0, d1, d2, d3, d7)>
+
+#encoding_nhwc_hwcf_input = #iree_encoding.encoding<operand_index = 0, op_type = conv,
+ element_types = [f32, f32, f32],
+ user_indexing_maps = [#nhwc_hwcf_map_in, #nhwc_hwcf_map_f, #nhwc_hwcf_map_out],
+ iteration_sizes = [1, 14, 14, 8, 3, 3, 4]>
+
+func.func @conv_input_pack(%arg0: tensor<1x16x16x4xf32>)
+ -> tensor<1x16x16x4xf32, #encoding_nhwc_hwcf_input>
+ attributes {
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz",
+ {target_triple="x86_64-xyz-xyz", cpu_features="+avx512f",
+ iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+} {
+ %0 = iree_encoding.set_encoding %arg0
+ : tensor<1x16x16x4xf32> -> tensor<1x16x16x4xf32, #encoding_nhwc_hwcf_input>
+ return %0 : tensor<1x16x16x4xf32, #encoding_nhwc_hwcf_input>
+}
+// CHECK-LABEL: func.func @conv_input_pack
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<1x16x16x4xf32>
+// CHECK: %[[PACK:.+]] = linalg.pack %[[ARG0]]
+// CHECK-SAME: outer_dims_perm = [0, 3, 1, 2]
+// CHECK-SAME: inner_dims_pos = [3]
+// CHECK-SAME: inner_tiles = [16]
+// CHECK-SAME: : tensor<1x16x16x4xf32> -> tensor<1x1x16x16x16xf32>
+// CHECK: return %[[PACK]]
+
+
+// Pack filter from [KH, KW, C, F] -> [OC/k0, IC/c0, KH, KW, c0, k0].
+
+#encoding_nhwc_hwcf_filter = #iree_encoding.encoding<operand_index = 1, op_type = conv,
+ element_types = [f32, f32, f32],
+ user_indexing_maps = [#nhwc_hwcf_map_in, #nhwc_hwcf_map_f, #nhwc_hwcf_map_out],
+ iteration_sizes = [1, 14, 14, 8, 3, 3, 4]>
+
+func.func @conv_filter_pack(%arg0: tensor<3x3x4x8xf32>)
+ -> tensor<3x3x4x8xf32, #encoding_nhwc_hwcf_filter>
+ attributes {
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz",
+ {target_triple="x86_64-xyz-xyz", cpu_features="+avx512f",
+ iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+} {
+ %0 = iree_encoding.set_encoding %arg0
+ : tensor<3x3x4x8xf32> -> tensor<3x3x4x8xf32, #encoding_nhwc_hwcf_filter>
+ return %0 : tensor<3x3x4x8xf32, #encoding_nhwc_hwcf_filter>
+}
+// CHECK-LABEL: func.func @conv_filter_pack
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<3x3x4x8xf32>
+// CHECK: %[[PACK:.+]] = linalg.pack %[[ARG0]]
+// CHECK-SAME: outer_dims_perm = [3, 2, 0, 1]
+// CHECK-SAME: inner_dims_pos = [2, 3]
+// CHECK-SAME: inner_tiles = [16, 16]
+// CHECK-SAME: : tensor<3x3x4x8xf32> -> tensor<1x1x3x3x16x16xf32>
+// CHECK: return %[[PACK]]
+
+
+// Unpack output from [N, OC/k0, OH, OW, k0] -> [N, OH, OW, OC].
+
+#encoding_nhwc_hwcf_output = #iree_encoding.encoding<operand_index = 2, op_type = conv,
+ element_types = [f32, f32, f32],
+ user_indexing_maps = [#nhwc_hwcf_map_in, #nhwc_hwcf_map_f, #nhwc_hwcf_map_out],
+ iteration_sizes = [1, 14, 14, 8, 3, 3, 4]>
+
+func.func @conv_output_unset(%arg0: tensor<1x14x14x8xf32, #encoding_nhwc_hwcf_output>)
+ -> tensor<1x14x14x8xf32>
+ attributes {
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz",
+ {target_triple="x86_64-xyz-xyz", cpu_features="+avx512f",
+ iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+} {
+ %0 = iree_encoding.unset_encoding %arg0
+ : tensor<1x14x14x8xf32, #encoding_nhwc_hwcf_output> -> tensor<1x14x14x8xf32>
+ return %0 : tensor<1x14x14x8xf32>
+}
+// CHECK-LABEL: func.func @conv_output_unset
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<1x1x14x14x16xf32>
+// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[ARG0]]
+// CHECK-SAME: outer_dims_perm = [0, 3, 1, 2]
+// CHECK-SAME: inner_dims_pos = [3]
+// CHECK-SAME: inner_tiles = [16]
+// CHECK-SAME: : tensor<1x1x14x14x16xf32> -> tensor<1x14x14x8xf32>
+// CHECK: return %[[UNPACK]]
+
+
+// Conv materialization.
+
+func.func @conv2d_nhwc_hwcf_materialize(
+ %input : tensor<1x16x16x4xf32, #encoding_nhwc_hwcf_input>,
+ %filter : tensor<3x3x4x8xf32, #encoding_nhwc_hwcf_filter>,
+ %output : tensor<1x14x14x8xf32, #encoding_nhwc_hwcf_output>)
+ -> tensor<1x14x14x8xf32, #encoding_nhwc_hwcf_output>
+ attributes {
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz",
+ {target_triple="x86_64-xyz-xyz", cpu_features="+avx512f",
+ iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+} {
+ %0 = linalg.conv_2d_nhwc_hwcf
+ {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
+ ins(%input, %filter
+ : tensor<1x16x16x4xf32, #encoding_nhwc_hwcf_input>,
+ tensor<3x3x4x8xf32, #encoding_nhwc_hwcf_filter>)
+ outs(%output : tensor<1x14x14x8xf32, #encoding_nhwc_hwcf_output>)
+ -> tensor<1x14x14x8xf32, #encoding_nhwc_hwcf_output>
+ return %0 : tensor<1x14x14x8xf32, #encoding_nhwc_hwcf_output>
+}
+// CHECK-LABEL: func.func @conv2d_nhwc_hwcf_materialize
+// CHECK-SAME: %[[INPUT:.+]]: tensor<1x1x16x16x16xf32>
+// CHECK-SAME: %[[FILTER:.+]]: tensor<1x1x3x3x16x16xf32>
+// CHECK-SAME: %[[OUTPUT:.+]]: tensor<1x1x14x14x16xf32>
+//
+// 9D data-tiled conv generic with three operands (input, filter, output):
+// CHECK: %[[RESULT:.+]] = linalg.generic
+// CHECK-SAME: indexing_maps = [#[[$M_IN]], #[[$M_FLT]], #[[$M_OUT]]]
+// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel",
+// CHECK-SAME: "reduction", "reduction", "reduction",
+// CHECK-SAME: "parallel", "reduction"]
+// CHECK-SAME: ins(%[[INPUT]], %[[FILTER]]
+// CHECK-SAME: outs(%[[OUTPUT]]
+// CHECK: %[[MUL:.+]] = arith.mulf
+// CHECK: %[[ADD:.+]] = arith.addf %[[MUL]]
+// CHECK: linalg.yield %[[ADD]]
+// CHECK: return %[[RESULT]]
+
+
+// Full conv materialization.
+
+func.func @conv2d_nhwc_hwcf_materialize_full(
+ %input : tensor<1x16x16x4xf32>,
+ %filter : tensor<3x3x4x8xf32>,
+ %output : tensor<1x14x14x8xf32>)
+ -> tensor<1x14x14x8xf32>
+ attributes {
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz",
+ {target_triple="x86_64-xyz-xyz", cpu_features="+avx512f",
+ iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+} {
+ %enc_in = iree_encoding.set_encoding %input
+ : tensor<1x16x16x4xf32> -> tensor<1x16x16x4xf32, #encoding_nhwc_hwcf_input>
+ %enc_f = iree_encoding.set_encoding %filter
+ : tensor<3x3x4x8xf32> -> tensor<3x3x4x8xf32, #encoding_nhwc_hwcf_filter>
+ %enc_out = iree_encoding.set_encoding %output
+ : tensor<1x14x14x8xf32> -> tensor<1x14x14x8xf32, #encoding_nhwc_hwcf_output>
+ %0 = linalg.conv_2d_nhwc_hwcf
+ {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
+ ins(%enc_in, %enc_f
+ : tensor<1x16x16x4xf32, #encoding_nhwc_hwcf_input>,
+ tensor<3x3x4x8xf32, #encoding_nhwc_hwcf_filter>)
+ outs(%enc_out : tensor<1x14x14x8xf32, #encoding_nhwc_hwcf_output>)
+ -> tensor<1x14x14x8xf32, #encoding_nhwc_hwcf_output>
+ %1 = iree_encoding.unset_encoding %0
+ : tensor<1x14x14x8xf32, #encoding_nhwc_hwcf_output> -> tensor<1x14x14x8xf32>
+ return %1 : tensor<1x14x14x8xf32>
+}
+// CHECK-LABEL: func.func @conv2d_nhwc_hwcf_materialize_full
+// CHECK-SAME: %[[IN:[a-zA-Z0-9]+]]: tensor<1x16x16x4xf32>
+// CHECK-SAME: %[[F:[a-zA-Z0-9]+]]: tensor<3x3x4x8xf32>
+// CHECK-SAME: %[[OUT:[a-zA-Z0-9]+]]: tensor<1x14x14x8xf32>
+//
+// CHECK: %[[PACK_IN:.+]] = linalg.pack %[[IN]]
+// CHECK-SAME: outer_dims_perm = [0, 3, 1, 2]
+// CHECK-SAME: inner_dims_pos = [3]
+// CHECK-SAME: inner_tiles = [16]
+// CHECK-SAME: : tensor<1x16x16x4xf32> -> tensor<1x1x16x16x16xf32>
+//
+// CHECK: %[[PACK_F:.+]] = linalg.pack %[[F]]
+// CHECK-SAME: outer_dims_perm = [3, 2, 0, 1]
+// CHECK-SAME: inner_dims_pos = [2, 3]
+// CHECK-SAME: inner_tiles = [16, 16]
+// CHECK-SAME: : tensor<3x3x4x8xf32> -> tensor<1x1x3x3x16x16xf32>
+//
+// CHECK: %[[PACK_OUT:.+]] = linalg.pack %[[OUT]]
+// CHECK-SAME: outer_dims_perm = [0, 3, 1, 2]
+// CHECK-SAME: inner_dims_pos = [3]
+// CHECK-SAME: inner_tiles = [16]
+// CHECK-SAME: : tensor<1x14x14x8xf32> -> tensor<1x1x14x14x16xf32>
+//
+// CHECK: %[[RESULT:.+]] = linalg.generic
+// CHECK-SAME: ins(%[[PACK_IN]], %[[PACK_F]]
+// CHECK-SAME: outs(%[[PACK_OUT]]
+//
+// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[RESULT]]
+// CHECK-SAME: outer_dims_perm = [0, 3, 1, 2]
+// CHECK-SAME: inner_dims_pos = [3]
+// CHECK-SAME: inner_tiles = [16]
+// CHECK-SAME: : tensor<1x1x14x14x16xf32> -> tensor<1x14x14x8xf32>
+
+
+// NHWC-HWCF conv expressed as a linalg.generic carrying conv encodings.
+// It must lower to the same pack / 9D generic / unpack as the named conv2d_nhwc_hwcf_materialize_full above.
+
+func.func @conv2d_nhwc_hwcf_materialize_generic(
+ %input : tensor<1x16x16x4xf32>,
+ %filter : tensor<3x3x4x8xf32>,
+ %output : tensor<1x14x14x8xf32>)
+ -> tensor<1x14x14x8xf32>
+ attributes {
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz",
+ {target_triple="x86_64-xyz-xyz", cpu_features="+avx512f",
+ iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+} {
+ %enc_in = iree_encoding.set_encoding %input
+ : tensor<1x16x16x4xf32> -> tensor<1x16x16x4xf32, #encoding_nhwc_hwcf_input>
+ %enc_f = iree_encoding.set_encoding %filter
+ : tensor<3x3x4x8xf32> -> tensor<3x3x4x8xf32, #encoding_nhwc_hwcf_filter>
+ %enc_out = iree_encoding.set_encoding %output
+ : tensor<1x14x14x8xf32> -> tensor<1x14x14x8xf32, #encoding_nhwc_hwcf_output>
+ %0 = linalg.generic
+ {indexing_maps = [#nhwc_hwcf_map_in, #nhwc_hwcf_map_f, #nhwc_hwcf_map_out],
+ iterator_types = ["parallel", "parallel", "parallel", "parallel",
+ "reduction", "reduction", "reduction"]}
+ ins(%enc_in, %enc_f
+ : tensor<1x16x16x4xf32, #encoding_nhwc_hwcf_input>,
+ tensor<3x3x4x8xf32, #encoding_nhwc_hwcf_filter>)
+ outs(%enc_out : tensor<1x14x14x8xf32, #encoding_nhwc_hwcf_output>) {
+ ^bb0(%in: f32, %f: f32, %acc: f32):
+ %m = arith.mulf %in, %f : f32
+ %a = arith.addf %acc, %m : f32
+ linalg.yield %a : f32
+ } -> tensor<1x14x14x8xf32, #encoding_nhwc_hwcf_output>
+ %1 = iree_encoding.unset_encoding %0
+ : tensor<1x14x14x8xf32, #encoding_nhwc_hwcf_output> -> tensor<1x14x14x8xf32>
+ return %1 : tensor<1x14x14x8xf32>
+}
+// CHECK-LABEL: func.func @conv2d_nhwc_hwcf_materialize_generic
+// CHECK-SAME: %[[IN:[a-zA-Z0-9]+]]: tensor<1x16x16x4xf32>
+// CHECK-SAME: %[[F:[a-zA-Z0-9]+]]: tensor<3x3x4x8xf32>
+// CHECK-SAME: %[[OUT:[a-zA-Z0-9]+]]: tensor<1x14x14x8xf32>
+//
+// CHECK: %[[PACK_IN:.+]] = linalg.pack %[[IN]]
+// CHECK-SAME: outer_dims_perm = [0, 3, 1, 2]
+// CHECK-SAME: inner_dims_pos = [3]
+// CHECK-SAME: inner_tiles = [16]
+// CHECK-SAME: : tensor<1x16x16x4xf32> -> tensor<1x1x16x16x16xf32>
+//
+// CHECK: %[[PACK_F:.+]] = linalg.pack %[[F]]
+// CHECK-SAME: outer_dims_perm = [3, 2, 0, 1]
+// CHECK-SAME: inner_dims_pos = [2, 3]
+// CHECK-SAME: inner_tiles = [16, 16]
+// CHECK-SAME: : tensor<3x3x4x8xf32> -> tensor<1x1x3x3x16x16xf32>
+//
+// CHECK: %[[PACK_OUT:.+]] = linalg.pack %[[OUT]]
+// CHECK-SAME: outer_dims_perm = [0, 3, 1, 2]
+// CHECK-SAME: inner_dims_pos = [3]
+// CHECK-SAME: inner_tiles = [16]
+// CHECK-SAME: : tensor<1x14x14x8xf32> -> tensor<1x1x14x14x16xf32>
+//
+// CHECK: %[[RESULT:.+]] = linalg.generic
+// CHECK-SAME: ins(%[[PACK_IN]], %[[PACK_F]]
+// CHECK-SAME: outs(%[[PACK_OUT]]
+//
+// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[RESULT]]
+// CHECK-SAME: outer_dims_perm = [0, 3, 1, 2]
+// CHECK-SAME: inner_dims_pos = [3]
+// CHECK-SAME: inner_tiles = [16]
+// CHECK-SAME: : tensor<1x1x14x14x16xf32> -> tensor<1x14x14x8xf32>
+
+
+// -----
+
+// conv_2d_nchw_fchw materialization.
+
+// Pack input from [N, C, H, W] -> [N, IC/c0, H, W, c0].
+
+#nchw_fchw_map_in = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d4, d2 + d5, d3 + d6)>
+#nchw_fchw_map_f = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d1, d4, d5, d6)>
+#nchw_fchw_map_out = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3)>
+
+// CHECK: #[[$M_IN:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d0, d4, d2 + d5, d3 + d6, d8)>
+// CHECK: #[[$M_FLT:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d1, d4, d5, d6, d8, d7)>
+// CHECK: #[[$M_OUT:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d0, d1, d2, d3, d7)>
+
+#encoding_nchw_fchw_input = #iree_encoding.encoding<operand_index = 0, op_type = conv,
+ element_types = [f32, f32, f32],
+ user_indexing_maps = [#nchw_fchw_map_in, #nchw_fchw_map_f, #nchw_fchw_map_out],
+ iteration_sizes = [1, 8, 14, 14, 4, 3, 3]>
+
+func.func @conv_nchw_input_pack(%arg0: tensor<1x4x16x16xf32>)
+ -> tensor<1x4x16x16xf32, #encoding_nchw_fchw_input>
+ attributes {
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz",
+ {target_triple="x86_64-xyz-xyz", cpu_features="+avx512f",
+ iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+} {
+ %0 = iree_encoding.set_encoding %arg0
+ : tensor<1x4x16x16xf32> -> tensor<1x4x16x16xf32, #encoding_nchw_fchw_input>
+ return %0 : tensor<1x4x16x16xf32, #encoding_nchw_fchw_input>
+}
+// CHECK-LABEL: func.func @conv_nchw_input_pack
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<1x4x16x16xf32>
+// CHECK: %[[PACK:.+]] = linalg.pack %[[ARG0]]
+// CHECK-SAME: outer_dims_perm = [0, 1, 2, 3]
+// CHECK-SAME: inner_dims_pos = [1]
+// CHECK-SAME: inner_tiles = [16]
+// CHECK-SAME: : tensor<1x4x16x16xf32> -> tensor<1x1x16x16x16xf32>
+// CHECK: return %[[PACK]]
+
+
+// Pack filter from [F, C, KH, KW] -> [OC/k0, IC/c0, KH, KW, c0, k0].
+
+#encoding_nchw_fchw_filter = #iree_encoding.encoding<operand_index = 1, op_type = conv,
+ element_types = [f32, f32, f32],
+ user_indexing_maps = [#nchw_fchw_map_in, #nchw_fchw_map_f, #nchw_fchw_map_out],
+ iteration_sizes = [1, 8, 14, 14, 4, 3, 3]>
+
+func.func @conv_fchw_filter_pack(%arg0: tensor<8x4x3x3xf32>)
+ -> tensor<8x4x3x3xf32, #encoding_nchw_fchw_filter>
+ attributes {
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz",
+ {target_triple="x86_64-xyz-xyz", cpu_features="+avx512f",
+ iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+} {
+ %0 = iree_encoding.set_encoding %arg0
+ : tensor<8x4x3x3xf32> -> tensor<8x4x3x3xf32, #encoding_nchw_fchw_filter>
+ return %0 : tensor<8x4x3x3xf32, #encoding_nchw_fchw_filter>
+}
+// CHECK-LABEL: func.func @conv_fchw_filter_pack
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<8x4x3x3xf32>
+// CHECK: %[[PACK:.+]] = linalg.pack %[[ARG0]]
+// CHECK-SAME: outer_dims_perm = [0, 1, 2, 3]
+// CHECK-SAME: inner_dims_pos = [1, 0]
+// CHECK-SAME: inner_tiles = [16, 16]
+// CHECK-SAME: : tensor<8x4x3x3xf32> -> tensor<1x1x3x3x16x16xf32>
+// CHECK: return %[[PACK]]
+
+
+// Unpack output from [N, OC/k0, OH, OW, k0] -> [N, OC, OH, OW].
+
+#encoding_nchw_fchw_output = #iree_encoding.encoding<operand_index = 2, op_type = conv,
+ element_types = [f32, f32, f32],
+ user_indexing_maps = [#nchw_fchw_map_in, #nchw_fchw_map_f, #nchw_fchw_map_out],
+ iteration_sizes = [1, 8, 14, 14, 4, 3, 3]>
+
+func.func @conv_nchw_output_unset(%arg0: tensor<1x8x14x14xf32, #encoding_nchw_fchw_output>)
+ -> tensor<1x8x14x14xf32>
+ attributes {
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz",
+ {target_triple="x86_64-xyz-xyz", cpu_features="+avx512f",
+ iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+} {
+ %0 = iree_encoding.unset_encoding %arg0
+ : tensor<1x8x14x14xf32, #encoding_nchw_fchw_output> -> tensor<1x8x14x14xf32>
+ return %0 : tensor<1x8x14x14xf32>
+}
+// CHECK-LABEL: func.func @conv_nchw_output_unset
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<1x1x14x14x16xf32>
+// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[ARG0]]
+// CHECK-SAME: outer_dims_perm = [0, 1, 2, 3]
+// CHECK-SAME: inner_dims_pos = [1]
+// CHECK-SAME: inner_tiles = [16]
+// CHECK-SAME: : tensor<1x1x14x14x16xf32> -> tensor<1x8x14x14xf32>
+// CHECK: return %[[UNPACK]]
+
+
+// Conv materialization.
+
+func.func @conv2d_nchw_fchw_materialize(
+ %input : tensor<1x4x16x16xf32, #encoding_nchw_fchw_input>,
+ %filter : tensor<8x4x3x3xf32, #encoding_nchw_fchw_filter>,
+ %output : tensor<1x8x14x14xf32, #encoding_nchw_fchw_output>)
+ -> tensor<1x8x14x14xf32, #encoding_nchw_fchw_output>
+ attributes {
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz",
+ {target_triple="x86_64-xyz-xyz", cpu_features="+avx512f",
+ iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+} {
+ %0 = linalg.conv_2d_nchw_fchw
+ {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
+ ins(%input, %filter
+ : tensor<1x4x16x16xf32, #encoding_nchw_fchw_input>,
+ tensor<8x4x3x3xf32, #encoding_nchw_fchw_filter>)
+ outs(%output : tensor<1x8x14x14xf32, #encoding_nchw_fchw_output>)
+ -> tensor<1x8x14x14xf32, #encoding_nchw_fchw_output>
+ return %0 : tensor<1x8x14x14xf32, #encoding_nchw_fchw_output>
+}
+// CHECK-LABEL: func.func @conv2d_nchw_fchw_materialize
+// CHECK-SAME: %[[INPUT:.+]]: tensor<1x1x16x16x16xf32>
+// CHECK-SAME: %[[FILTER:.+]]: tensor<1x1x3x3x16x16xf32>
+// CHECK-SAME: %[[OUTPUT:.+]]: tensor<1x1x14x14x16xf32>
+//
+// CHECK: %[[RESULT:.+]] = linalg.generic
+// CHECK-SAME: indexing_maps = [#[[$M_IN]], #[[$M_FLT]], #[[$M_OUT]]]
+// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel",
+// CHECK-SAME: "reduction", "reduction", "reduction",
+// CHECK-SAME: "parallel", "reduction"]
+// CHECK-SAME: ins(%[[INPUT]], %[[FILTER]]
+// CHECK-SAME: outs(%[[OUTPUT]]
+// CHECK: %[[MUL:.+]] = arith.mulf
+// CHECK: %[[ADD:.+]] = arith.addf %[[MUL]]
+// CHECK: linalg.yield %[[ADD]]
+// CHECK: return %[[RESULT]]
+
+
+// Full conv materialization.
+
+func.func @conv2d_nchw_fchw_materialize_full(
+ %input : tensor<1x4x16x16xf32>,
+ %filter : tensor<8x4x3x3xf32>,
+ %output : tensor<1x8x14x14xf32>)
+ -> tensor<1x8x14x14xf32>
+ attributes {
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz",
+ {target_triple="x86_64-xyz-xyz", cpu_features="+avx512f",
+ iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+} {
+ %enc_in = iree_encoding.set_encoding %input
+ : tensor<1x4x16x16xf32> -> tensor<1x4x16x16xf32, #encoding_nchw_fchw_input>
+ %enc_f = iree_encoding.set_encoding %filter
+ : tensor<8x4x3x3xf32> -> tensor<8x4x3x3xf32, #encoding_nchw_fchw_filter>
+ %enc_out = iree_encoding.set_encoding %output
+ : tensor<1x8x14x14xf32> -> tensor<1x8x14x14xf32, #encoding_nchw_fchw_output>
+ %0 = linalg.conv_2d_nchw_fchw
+ {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
+ ins(%enc_in, %enc_f
+ : tensor<1x4x16x16xf32, #encoding_nchw_fchw_input>,
+ tensor<8x4x3x3xf32, #encoding_nchw_fchw_filter>)
+ outs(%enc_out : tensor<1x8x14x14xf32, #encoding_nchw_fchw_output>)
+ -> tensor<1x8x14x14xf32, #encoding_nchw_fchw_output>
+ %1 = iree_encoding.unset_encoding %0
+ : tensor<1x8x14x14xf32, #encoding_nchw_fchw_output> -> tensor<1x8x14x14xf32>
+ return %1 : tensor<1x8x14x14xf32>
+}
+// CHECK-LABEL: func.func @conv2d_nchw_fchw_materialize_full
+// CHECK-SAME: %[[IN:[a-zA-Z0-9]+]]: tensor<1x4x16x16xf32>
+// CHECK-SAME: %[[F:[a-zA-Z0-9]+]]: tensor<8x4x3x3xf32>
+// CHECK-SAME: %[[OUT:[a-zA-Z0-9]+]]: tensor<1x8x14x14xf32>
+//
+// CHECK: %[[PACK_IN:.+]] = linalg.pack %[[IN]]
+// CHECK-SAME: outer_dims_perm = [0, 1, 2, 3]
+// CHECK-SAME: inner_dims_pos = [1]
+// CHECK-SAME: inner_tiles = [16]
+// CHECK-SAME: : tensor<1x4x16x16xf32> -> tensor<1x1x16x16x16xf32>
+//
+// CHECK: %[[PACK_F:.+]] = linalg.pack %[[F]]
+// CHECK-SAME: outer_dims_perm = [0, 1, 2, 3]
+// CHECK-SAME: inner_dims_pos = [1, 0]
+// CHECK-SAME: inner_tiles = [16, 16]
+// CHECK-SAME: : tensor<8x4x3x3xf32> -> tensor<1x1x3x3x16x16xf32>
+//
+// CHECK: %[[PACK_OUT:.+]] = linalg.pack %[[OUT]]
+// CHECK-SAME: outer_dims_perm = [0, 1, 2, 3]
+// CHECK-SAME: inner_dims_pos = [1]
+// CHECK-SAME: inner_tiles = [16]
+// CHECK-SAME: : tensor<1x8x14x14xf32> -> tensor<1x1x14x14x16xf32>
+//
+// CHECK: %[[RESULT:.+]] = linalg.generic
+// CHECK-SAME: ins(%[[PACK_IN]], %[[PACK_F]]
+// CHECK-SAME: outs(%[[PACK_OUT]]
+//
+// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[RESULT]]
+// CHECK-SAME: outer_dims_perm = [0, 1, 2, 3]
+// CHECK-SAME: inner_dims_pos = [1]
+// CHECK-SAME: inner_tiles = [16]
+// CHECK-SAME: : tensor<1x1x14x14x16xf32> -> tensor<1x8x14x14xf32>
+
+// -----
+
+// conv_2d_nhwc_fhwc materialization.
+
+// Pack filter from [F, KH, KW, C] -> [OC/k0, IC/c0, KH, KW, c0, k0].
+
+#nhwc_fhwc_map_in = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1 + d4, d2 + d5, d6)>
+#nhwc_fhwc_map_f = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d3, d4, d5, d6)>
+#nhwc_fhwc_map_out = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3)>
+
+#encoding_nhwc_fhwc_filter = #iree_encoding.encoding<operand_index = 1, op_type = conv,
+ element_types = [f32, f32, f32],
+ user_indexing_maps = [#nhwc_fhwc_map_in, #nhwc_fhwc_map_f, #nhwc_fhwc_map_out],
+ iteration_sizes = [1, 14, 14, 8, 3, 3, 4]>
+
+func.func @conv_fhwc_filter_pack(%arg0: tensor<8x3x3x4xf32>)
+ -> tensor<8x3x3x4xf32, #encoding_nhwc_fhwc_filter>
+ attributes {
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz",
+ {target_triple="x86_64-xyz-xyz", cpu_features="+avx512f",
+ iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+} {
+ %0 = iree_encoding.set_encoding %arg0
+ : tensor<8x3x3x4xf32> -> tensor<8x3x3x4xf32, #encoding_nhwc_fhwc_filter>
+ return %0 : tensor<8x3x3x4xf32, #encoding_nhwc_fhwc_filter>
+}
+// CHECK-LABEL: func.func @conv_fhwc_filter_pack
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<8x3x3x4xf32>
+// CHECK: %[[PACK:.+]] = linalg.pack %[[ARG0]]
+// CHECK-SAME: outer_dims_perm = [0, 3, 1, 2]
+// CHECK-SAME: inner_dims_pos = [3, 0]
+// CHECK-SAME: inner_tiles = [16, 16]
+// CHECK-SAME: : tensor<8x3x3x4xf32> -> tensor<1x1x3x3x16x16xf32>
+// CHECK: return %[[PACK]]
+
+// -----
+
+// Unit batch-dim folded conv: N=1 stripped upstream, so operands pack to rank-4
+// and the lowering restores the batch via expand/collapse around the 9-D generic.
+
+#bl_map_in = affine_map<(oh, ow, oc, fh, fw, ic) -> (oh + fh, ow + fw, ic)>
+#bl_map_f = affine_map<(oh, ow, oc, fh, fw, ic) -> (fh, fw, ic, oc)>
+#bl_map_out = affine_map<(oh, ow, oc, fh, fw, ic) -> (oh, ow, oc)>
+
+#encoding_bl_input = #iree_encoding.encoding<operand_index = 0, op_type = conv,
+ element_types = [f32, f32, f32],
+ user_indexing_maps = [#bl_map_in, #bl_map_f, #bl_map_out],
+ iteration_sizes = [14, 14, 8, 3, 3, 4]>
+#encoding_bl_filter = #iree_encoding.encoding<operand_index = 1, op_type = conv,
+ element_types = [f32, f32, f32],
+ user_indexing_maps = [#bl_map_in, #bl_map_f, #bl_map_out],
+ iteration_sizes = [14, 14, 8, 3, 3, 4]>
+#encoding_bl_output = #iree_encoding.encoding<operand_index = 2, op_type = conv,
+ element_types = [f32, f32, f32],
+ user_indexing_maps = [#bl_map_in, #bl_map_f, #bl_map_out],
+ iteration_sizes = [14, 14, 8, 3, 3, 4]>
+
+func.func @conv2d_batchless_materialize(
+ %input : tensor<16x16x4xf32, #encoding_bl_input>,
+ %filter : tensor<3x3x4x8xf32, #encoding_bl_filter>,
+ %output : tensor<14x14x8xf32, #encoding_bl_output>)
+ -> tensor<14x14x8xf32, #encoding_bl_output>
+ attributes {
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz",
+ {target_triple="x86_64-xyz-xyz", cpu_features="+avx512f",
+ iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+} {
+ %0 = linalg.generic
+ {indexing_maps = [#bl_map_in, #bl_map_f, #bl_map_out],
+ iterator_types = ["parallel", "parallel", "parallel",
+ "reduction", "reduction", "reduction"]}
+ ins(%input, %filter
+ : tensor<16x16x4xf32, #encoding_bl_input>,
+ tensor<3x3x4x8xf32, #encoding_bl_filter>)
+ outs(%output : tensor<14x14x8xf32, #encoding_bl_output>) {
+ ^bb0(%in: f32, %f: f32, %acc: f32):
+ %m = arith.mulf %in, %f : f32
+ %a = arith.addf %acc, %m : f32
+ linalg.yield %a : f32
+ } -> tensor<14x14x8xf32, #encoding_bl_output>
+ return %0 : tensor<14x14x8xf32, #encoding_bl_output>
+}
+// CHECK-LABEL: func.func @conv2d_batchless_materialize
+// CHECK-SAME: %[[INPUT:.+]]: tensor<1x16x16x16xf32>
+// CHECK-SAME: %[[FILTER:.+]]: tensor<1x1x3x3x16x16xf32>
+// CHECK-SAME: %[[OUTPUT:.+]]: tensor<1x14x14x16xf32>
+// CHECK: %[[EXP_IN:.+]] = tensor.expand_shape %[[INPUT]] {{\[}}[0, 1], [2], [3], [4]{{\]}}
+// CHECK-SAME: : tensor<1x16x16x16xf32> into tensor<1x1x16x16x16xf32>
+// CHECK: %[[EXP_OUT:.+]] = tensor.expand_shape %[[OUTPUT]] {{\[}}[0, 1], [2], [3], [4]{{\]}}
+// CHECK-SAME: : tensor<1x14x14x16xf32> into tensor<1x1x14x14x16xf32>
+// CHECK: %[[RESULT:.+]] = linalg.generic
+// CHECK-SAME: indexing_maps = [#[[$M_IN]], #[[$M_FLT]], #[[$M_OUT]]]
+// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel",
+// CHECK-SAME: "reduction", "reduction", "reduction",
+// CHECK-SAME: "parallel", "reduction"]
+// CHECK-SAME: ins(%[[EXP_IN]], %[[FILTER]]
+// CHECK-SAME: outs(%[[EXP_OUT]]
+// CHECK: %[[MUL:.+]] = arith.mulf
+// CHECK: %[[ADD:.+]] = arith.addf %[[MUL]]
+// CHECK: linalg.yield %[[ADD]]
+// CHECK: %[[COLLAPSE:.+]] = tensor.collapse_shape %[[RESULT]] {{\[}}[0, 1], [2], [3], [4]{{\]}}
+// CHECK-SAME: : tensor<1x1x14x14x16xf32> into tensor<1x14x14x16xf32>
+// CHECK: return %[[COLLAPSE]]
+
+// -----
+
+// Unsupported convs: the resolver only data-tiles plain 2D convs, so grouped,
+// non-2D (1D/3D), and dilated convs resolve to identity layout.
+
+// 1D conv (nwc_wcf)
+
+#map_1d_in = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1 + d3, d4)>
+#map_1d_flt = affine_map<(d0, d1, d2, d3, d4) -> (d3, d4, d2)>
+#map_1d_out = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2)>
+#encoding_1d_in = #iree_encoding.encoding<operand_index = 0, op_type = conv,
+ element_types = [f32, f32, f32],
+ user_indexing_maps = [#map_1d_in, #map_1d_flt, #map_1d_out],
+ iteration_sizes = [1, 14, 8, 3, 4]>
+#encoding_1d_flt = #iree_encoding.encoding<operand_index = 1, op_type = conv,
+ element_types = [f32, f32, f32],
+ user_indexing_maps = [#map_1d_in, #map_1d_flt, #map_1d_out],
+ iteration_sizes = [1, 14, 8, 3, 4]>
+#encoding_1d_out = #iree_encoding.encoding<operand_index = 2, op_type = conv,
+ element_types = [f32, f32, f32],
+ user_indexing_maps = [#map_1d_in, #map_1d_flt, #map_1d_out],
+ iteration_sizes = [1, 14, 8, 3, 4]>
+func.func @conv1d_unsupported_identity(
+ %in : tensor<1x16x4xf32>,
+ %flt : tensor<3x4x8xf32>,
+ %out : tensor<1x14x8xf32>)
+ -> tensor<1x14x8xf32>
+ attributes {
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz",
+ {target_triple="x86_64-xyz-xyz", cpu_features="+avx512f",
+ iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+} {
+ %e_in = iree_encoding.set_encoding %in
+ : tensor<1x16x4xf32> -> tensor<1x16x4xf32, #encoding_1d_in>
+ %e_flt = iree_encoding.set_encoding %flt
+ : tensor<3x4x8xf32> -> tensor<3x4x8xf32, #encoding_1d_flt>
+ %e_out = iree_encoding.set_encoding %out
+ : tensor<1x14x8xf32> -> tensor<1x14x8xf32, #encoding_1d_out>
+ %0 = linalg.conv_1d_nwc_wcf
+ {dilations = dense<1> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>}
+ ins(%e_in, %e_flt
+ : tensor<1x16x4xf32, #encoding_1d_in>, tensor<3x4x8xf32, #encoding_1d_flt>)
+ outs(%e_out : tensor<1x14x8xf32, #encoding_1d_out>)
+ -> tensor<1x14x8xf32, #encoding_1d_out>
+ %1 = iree_encoding.unset_encoding %0
+ : tensor<1x14x8xf32, #encoding_1d_out> -> tensor<1x14x8xf32>
+ return %1 : tensor<1x14x8xf32>
+}
+// CHECK-LABEL: func.func @conv1d_unsupported_identity
+// CHECK-SAME: %[[IN:[a-zA-Z0-9]+]]: tensor<1x16x4xf32>
+// CHECK-SAME: %[[FLT:[a-zA-Z0-9]+]]: tensor<3x4x8xf32>
+// CHECK-SAME: %[[OUT:[a-zA-Z0-9]+]]: tensor<1x14x8xf32>
+// CHECK-NOT: linalg.pack
+// CHECK-NOT: iree_encoding
+// CHECK: %[[CONV:.+]] = linalg.conv_1d_nwc_wcf
+// CHECK-SAME: ins(%[[IN]], %[[FLT]] : tensor<1x16x4xf32>, tensor<3x4x8xf32>)
+// CHECK-SAME: outs(%[[OUT]] : tensor<1x14x8xf32>)
+// CHECK-NOT: linalg.unpack
+// CHECK: return %[[CONV]]
+
+// -----
+
+// 2D grouped conv (ngchw_fgchw)
+
+#map_g_in = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d0, d1, d5, d3 + d6, d4 + d7)>
+#map_g_flt = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d2, d1, d5, d6, d7)>
+#map_g_out = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d0, d1, d2, d3, d4)>
+#encoding_g_in = #iree_encoding.encoding<operand_index = 0, op_type = conv,
+ element_types = [f32, f32, f32],
+ user_indexing_maps = [#map_g_in, #map_g_flt, #map_g_out],
+ iteration_sizes = [1, 2, 8, 14, 14, 4, 3, 3]>
+#encoding_g_flt = #iree_encoding.encoding<operand_index = 1, op_type = conv,
+ element_types = [f32, f32, f32],
+ user_indexing_maps = [#map_g_in, #map_g_flt, #map_g_out],
+ iteration_sizes = [1, 2, 8, 14, 14, 4, 3, 3]>
+#encoding_g_out = #iree_encoding.encoding<operand_index = 2, op_type = conv,
+ element_types = [f32, f32, f32],
+ user_indexing_maps = [#map_g_in, #map_g_flt, #map_g_out],
+ iteration_sizes = [1, 2, 8, 14, 14, 4, 3, 3]>
+func.func @conv2d_grouped_unsupported_identity(
+ %in : tensor<1x2x4x16x16xf32>,
+ %flt : tensor<8x2x4x3x3xf32>,
+ %out : tensor<1x2x8x14x14xf32>)
+ -> tensor<1x2x8x14x14xf32>
+ attributes {
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz",
+ {target_triple="x86_64-xyz-xyz", cpu_features="+avx512f",
+ iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+} {
+ %e_in = iree_encoding.set_encoding %in
+ : tensor<1x2x4x16x16xf32> -> tensor<1x2x4x16x16xf32, #encoding_g_in>
+ %e_flt = iree_encoding.set_encoding %flt
+ : tensor<8x2x4x3x3xf32> -> tensor<8x2x4x3x3xf32, #encoding_g_flt>
+ %e_out = iree_encoding.set_encoding %out
+ : tensor<1x2x8x14x14xf32> -> tensor<1x2x8x14x14xf32, #encoding_g_out>
+ %0 = linalg.conv_2d_ngchw_fgchw
+ {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
+ ins(%e_in, %e_flt
+ : tensor<1x2x4x16x16xf32, #encoding_g_in>, tensor<8x2x4x3x3xf32, #encoding_g_flt>)
+ outs(%e_out : tensor<1x2x8x14x14xf32, #encoding_g_out>)
+ -> tensor<1x2x8x14x14xf32, #encoding_g_out>
+ %1 = iree_encoding.unset_encoding %0
+ : tensor<1x2x8x14x14xf32, #encoding_g_out> -> tensor<1x2x8x14x14xf32>
+ return %1 : tensor<1x2x8x14x14xf32>
+}
+// CHECK-LABEL: func.func @conv2d_grouped_unsupported_identity
+// CHECK-SAME: %[[IN:[a-zA-Z0-9]+]]: tensor<1x2x4x16x16xf32>
+// CHECK-SAME: %[[FLT:[a-zA-Z0-9]+]]: tensor<8x2x4x3x3xf32>
+// CHECK-SAME: %[[OUT:[a-zA-Z0-9]+]]: tensor<1x2x8x14x14xf32>
+// CHECK-NOT: linalg.pack
+// CHECK-NOT: iree_encoding
+// CHECK: %[[CONV:.+]] = linalg.conv_2d_ngchw_fgchw
+// CHECK-SAME: ins(%[[IN]], %[[FLT]] : tensor<1x2x4x16x16xf32>, tensor<8x2x4x3x3xf32>)
+// CHECK-SAME: outs(%[[OUT]] : tensor<1x2x8x14x14xf32>)
+// CHECK-NOT: linalg.unpack
+// CHECK: return %[[CONV]]
+
+// -----
+
+// 3D conv (ndhwc_dhwcf)
+
+#map_3d_in = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d0, d1 + d5, d2 + d6, d3 + d7, d8)>
+#map_3d_flt = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d5, d6, d7, d8, d4)>
+#map_3d_out = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d0, d1, d2, d3, d4)>
+#encoding_3d_in = #iree_encoding.encoding<operand_index = 0, op_type = conv,
+ element_types = [f32, f32, f32],
+ user_indexing_maps = [#map_3d_in, #map_3d_flt, #map_3d_out],
+ iteration_sizes = [1, 14, 14, 14, 8, 3, 3, 3, 4]>
+#encoding_3d_flt = #iree_encoding.encoding<operand_index = 1, op_type = conv,
+ element_types = [f32, f32, f32],
+ user_indexing_maps = [#map_3d_in, #map_3d_flt, #map_3d_out],
+ iteration_sizes = [1, 14, 14, 14, 8, 3, 3, 3, 4]>
+#encoding_3d_out = #iree_encoding.encoding<operand_index = 2, op_type = conv,
+ element_types = [f32, f32, f32],
+ user_indexing_maps = [#map_3d_in, #map_3d_flt, #map_3d_out],
+ iteration_sizes = [1, 14, 14, 14, 8, 3, 3, 3, 4]>
+func.func @conv3d_unsupported_identity(
+ %in : tensor<1x16x16x16x4xf32>,
+ %flt : tensor<3x3x3x4x8xf32>,
+ %out : tensor<1x14x14x14x8xf32>)
+ -> tensor<1x14x14x14x8xf32>
+ attributes {
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz",
+ {target_triple="x86_64-xyz-xyz", cpu_features="+avx512f",
+ iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+} {
+ %e_in = iree_encoding.set_encoding %in
+ : tensor<1x16x16x16x4xf32> -> tensor<1x16x16x16x4xf32, #encoding_3d_in>
+ %e_flt = iree_encoding.set_encoding %flt
+ : tensor<3x3x3x4x8xf32> -> tensor<3x3x3x4x8xf32, #encoding_3d_flt>
+ %e_out = iree_encoding.set_encoding %out
+ : tensor<1x14x14x14x8xf32> -> tensor<1x14x14x14x8xf32, #encoding_3d_out>
+ %0 = linalg.conv_3d_ndhwc_dhwcf
+ {dilations = dense<1> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>}
+ ins(%e_in, %e_flt
+ : tensor<1x16x16x16x4xf32, #encoding_3d_in>, tensor<3x3x3x4x8xf32, #encoding_3d_flt>)
+ outs(%e_out : tensor<1x14x14x14x8xf32, #encoding_3d_out>)
+ -> tensor<1x14x14x14x8xf32, #encoding_3d_out>
+ %1 = iree_encoding.unset_encoding %0
+ : tensor<1x14x14x14x8xf32, #encoding_3d_out> -> tensor<1x14x14x14x8xf32>
+ return %1 : tensor<1x14x14x14x8xf32>
+}
+// CHECK-LABEL: func.func @conv3d_unsupported_identity
+// CHECK-SAME: %[[IN:[a-zA-Z0-9]+]]: tensor<1x16x16x16x4xf32>
+// CHECK-SAME: %[[FLT:[a-zA-Z0-9]+]]: tensor<3x3x3x4x8xf32>
+// CHECK-SAME: %[[OUT:[a-zA-Z0-9]+]]: tensor<1x14x14x14x8xf32>
+// CHECK-NOT: linalg.pack
+// CHECK-NOT: iree_encoding
+// CHECK: %[[CONV:.+]] = linalg.conv_3d_ndhwc_dhwcf
+// CHECK-SAME: ins(%[[IN]], %[[FLT]] : tensor<1x16x16x16x4xf32>, tensor<3x3x3x4x8xf32>)
+// CHECK-SAME: outs(%[[OUT]] : tensor<1x14x14x14x8xf32>)
+// CHECK-NOT: linalg.unpack
+// CHECK: return %[[CONV]]
+
+// -----
+
+// Dilated 2D conv
+
+#map_d_in = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1 + d4 * 2, d2 + d5 * 2, d6)>
+#map_d_flt = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d4, d5, d6, d3)>
+#map_d_out = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3)>
+#encoding_d_in = #iree_encoding.encoding<operand_index = 0, op_type = conv,
+ element_types = [f32, f32, f32],
+ user_indexing_maps = [#map_d_in, #map_d_flt, #map_d_out],
+ iteration_sizes = [1, 12, 12, 8, 3, 3, 4]>
+#encoding_d_flt = #iree_encoding.encoding<operand_index = 1, op_type = conv,
+ element_types = [f32, f32, f32],
+ user_indexing_maps = [#map_d_in, #map_d_flt, #map_d_out],
+ iteration_sizes = [1, 12, 12, 8, 3, 3, 4]>
+#encoding_d_out = #iree_encoding.encoding<operand_index = 2, op_type = conv,
+ element_types = [f32, f32, f32],
+ user_indexing_maps = [#map_d_in, #map_d_flt, #map_d_out],
+ iteration_sizes = [1, 12, 12, 8, 3, 3, 4]>
+func.func @conv2d_dilated_unsupported_identity(
+ %in : tensor<1x16x16x4xf32>,
+ %flt : tensor<3x3x4x8xf32>,
+ %out : tensor<1x12x12x8xf32>)
+ -> tensor<1x12x12x8xf32>
+ attributes {
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz",
+ {target_triple="x86_64-xyz-xyz", cpu_features="+avx512f",
+ iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+} {
+ %e_in = iree_encoding.set_encoding %in
+ : tensor<1x16x16x4xf32> -> tensor<1x16x16x4xf32, #encoding_d_in>
+ %e_flt = iree_encoding.set_encoding %flt
+ : tensor<3x3x4x8xf32> -> tensor<3x3x4x8xf32, #encoding_d_flt>
+ %e_out = iree_encoding.set_encoding %out
+ : tensor<1x12x12x8xf32> -> tensor<1x12x12x8xf32, #encoding_d_out>
+ %0 = linalg.conv_2d_nhwc_hwcf
+ {dilations = dense<2> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
+ ins(%e_in, %e_flt
+ : tensor<1x16x16x4xf32, #encoding_d_in>, tensor<3x3x4x8xf32, #encoding_d_flt>)
+ outs(%e_out : tensor<1x12x12x8xf32, #encoding_d_out>)
+ -> tensor<1x12x12x8xf32, #encoding_d_out>
+ %1 = iree_encoding.unset_encoding %0
+ : tensor<1x12x12x8xf32, #encoding_d_out> -> tensor<1x12x12x8xf32>
+ return %1 : tensor<1x12x12x8xf32>
+}
+// CHECK-LABEL: func.func @conv2d_dilated_unsupported_identity
+// CHECK-SAME: %[[IN:[a-zA-Z0-9]+]]: tensor<1x16x16x4xf32>
+// CHECK-SAME: %[[FLT:[a-zA-Z0-9]+]]: tensor<3x3x4x8xf32>
+// CHECK-SAME: %[[OUT:[a-zA-Z0-9]+]]: tensor<1x12x12x8xf32>
+// CHECK-NOT: linalg.pack
+// CHECK-NOT: iree_encoding
+// CHECK: %[[CONV:.+]] = linalg.conv_2d_nhwc_hwcf
+// CHECK-SAME: {dilations = dense<2> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
+// CHECK-SAME: ins(%[[IN]], %[[FLT]] : tensor<1x16x16x4xf32>, tensor<3x3x4x8xf32>)
+// CHECK-SAME: outs(%[[OUT]] : tensor<1x12x12x8xf32>)
+// CHECK-NOT: linalg.unpack
+// CHECK: return %[[CONV]]
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils/BUILD.bazel b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils/BUILD.bazel
index 70299ac..724165f 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils/BUILD.bazel
+++ b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils/BUILD.bazel
@@ -30,6 +30,7 @@
"@llvm-project//mlir:DialectUtils",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:LinalgDialect",
+ "@llvm-project//mlir:LinalgInterfaces",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TensorDialect",
],
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils/CMakeLists.txt
index 1870c9f..5f4ce3a 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils/CMakeLists.txt
@@ -22,6 +22,7 @@
MLIRArithDialect
MLIRIR
MLIRLinalgDialect
+ MLIRLinalgInterfacesIncGenLib
MLIRSupport
MLIRTensorDialect
iree::compiler::Codegen::Dialect::Codegen::IR::IREECodegenDialect
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 beb9b59..d49aad4 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.cpp
+++ b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.cpp
@@ -10,7 +10,11 @@
#include "iree/compiler/Dialect/LinalgExt/Utils/MatchUtils.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/InterleavedRange.h"
+#include "mlir/Dialect/Linalg/IR/LinalgInterfaces.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
+#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
+#include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/AffineMap.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinAttributes.h"
@@ -526,4 +530,164 @@
return encodingInfo;
}
+FailureOr<MaterializeEncodingInfo>
+getEncodingInfoForConv(Encoding::EncodingAttr encoding, TileOCxIC tile) {
+ int64_t operandIdx = encoding.getOperandIndex().getInt();
+ SmallVector<AffineMap> maps = encoding.getRootMaps();
+
+ FailureOr<linalg::ConvolutionDimensions> cDims =
+ linalg::inferConvolutionDims(maps);
+ if (failed(cDims)) {
+ return failure();
+ }
+
+ // Data tiling requires both input and output channel dimensions.
+ if (cDims->inputChannel.empty() || cDims->outputChannel.empty()) {
+ return failure();
+ }
+
+ // Grouped/depthwise convolutions are not supported. A depth dimension
+ // appears as a bare dim in the input map and would be misclassified as
+ // an output image dimension by the position inference below.
+ if (!cDims->depth.empty()) {
+ return failure();
+ }
+
+ // Map classified loop dims to tensor positions for this operand.
+ //
+ // A loop dimension can be mapped to a tensor position only if it appears
+ // as a bare dimension in the operand's indexing map. Dimensions that occur
+ // only inside a compound affine expression have no single tensor position
+ // and cannot be recovered this way.
+ //
+ // For example: linalg.conv_2d_nhwc_hwcf has input indexing map
+ // (n, oh + fh, ow + fw, ic)
+ // where the loop dimensions are
+ // (n, oh, ow, oc, fh, fw, ic) = (0, 1, 2, 3, 4, 5, 6).
+ // Only n and ic appear as bare dimensions, at tensor positions 0 and 3.
+ // The outputImage dimensions (oh, ow) appear only in compound expressions, so
+ // they cannot be recovered here.
+ auto collectBarePositions =
+ [&encoding](ArrayRef<unsigned> loopDims) -> SmallVector<int64_t> {
+ SmallVector<int64_t> positions;
+ for (unsigned d : loopDims) {
+ if (auto pos = encoding.mapDimToOperandIndex(d)) {
+ positions.push_back(static_cast<int64_t>(*pos));
+ }
+ }
+ return positions;
+ };
+
+ // For the input operand, recover the outputImage tensor positions.
+ //
+ // From the example above, the known positions are {0, 3}, so the remaining
+ // tensor positions {1, 2} correspond to the outputImage dimensions.
+ auto collectRemainingPositions =
+ [](int64_t rank,
+ ArrayRef<int64_t> knownPositions) -> SmallVector<int64_t> {
+ llvm::SmallDenseSet<int64_t> known(knownPositions.begin(),
+ knownPositions.end());
+ SmallVector<int64_t> remaining;
+ for (int64_t i = 0; i < rank; ++i) {
+ if (!known.contains(i)) {
+ remaining.push_back(i);
+ }
+ }
+ return remaining;
+ };
+
+ SmallVector<int64_t> batchPos = collectBarePositions(cDims->batch);
+ SmallVector<int64_t> ocPos = collectBarePositions(cDims->outputChannel);
+ SmallVector<int64_t> flPos = collectBarePositions(cDims->filterLoop);
+ SmallVector<int64_t> icPos = collectBarePositions(cDims->inputChannel);
+
+ MaterializeEncodingInfo info;
+
+ if (operandIdx == Encoding::CONV_IN) {
+ if (icPos.empty()) {
+ return failure();
+ }
+ int64_t rank = maps[operandIdx].getNumResults();
+ SmallVector<int64_t> knownInputPos;
+ knownInputPos.append(batchPos);
+ knownInputPos.append(icPos);
+ SmallVector<int64_t> outputImagePosInInput =
+ collectRemainingPositions(rank, knownInputPos);
+
+ SmallVector<int64_t> canonicalOrder;
+ canonicalOrder.append(batchPos);
+ canonicalOrder.append(icPos);
+ canonicalOrder.append(outputImagePosInInput);
+
+ info.outerDimsPerm = canonicalOrder;
+ info.innerDimsPos = {icPos[0]};
+ info.innerTileSizes = {tile.IC};
+ } else if (operandIdx == Encoding::CONV_FILTER) {
+ if (ocPos.empty() || icPos.empty()) {
+ return failure();
+ }
+ SmallVector<int64_t> canonicalOrder;
+ canonicalOrder.append(ocPos);
+ canonicalOrder.append(icPos);
+ canonicalOrder.append(flPos);
+
+ info.outerDimsPerm = canonicalOrder;
+ info.innerDimsPos = {icPos[0], ocPos[0]};
+ info.innerTileSizes = {tile.IC, tile.OC};
+ } else if (operandIdx == Encoding::CONV_OUT) {
+ if (ocPos.empty()) {
+ return failure();
+ }
+ SmallVector<int64_t> outputImagePosInOutput =
+ collectBarePositions(cDims->outputImage);
+
+ SmallVector<int64_t> canonicalOrder;
+ canonicalOrder.append(batchPos);
+ canonicalOrder.append(ocPos);
+ canonicalOrder.append(outputImagePosInOutput);
+
+ info.outerDimsPerm = canonicalOrder;
+ info.innerDimsPos = {ocPos[0]};
+ info.innerTileSizes = {tile.OC};
+ } else {
+ return failure();
+ }
+
+ return info;
+}
+
+DataTiledConvIterationSpace
+getDataTiledConvIterationSpace(MLIRContext *ctx, ArrayRef<int64_t> strides,
+ ArrayRef<int64_t> dilations) {
+ assert(strides.size() == 2 && dilations.size() == 2 &&
+ "expected 2D convolution strides/dilations");
+ AffineExpr e0, e1, e2, e3, e4, e5, e6, e7, e8;
+ bindDims(ctx, e0, e1, e2, e3, e4, e5, e6, e7, e8);
+
+ DataTiledConvIterationSpace space;
+ space.indexingMaps = {
+ // input: [N, IC/c0, H, W, c0]
+ AffineMap::get(9, 0,
+ {e0, e4, e2 * strides[0] + e5 * dilations[0],
+ e3 * strides[1] + e6 * dilations[1], e8},
+ ctx),
+ // filter: [OC/k0, IC/c0, FH, FW, c0, k0]
+ AffineMap::get(9, 0, {e1, e4, e5, e6, e8, e7}, ctx),
+ // output: [N, OC/k0, OH, OW, k0]
+ AffineMap::get(9, 0, {e0, e1, e2, e3, e7}, ctx),
+ };
+ space.iteratorTypes = {
+ utils::IteratorType::parallel, // d0 = n
+ utils::IteratorType::parallel, // d1 = oc_outer
+ utils::IteratorType::parallel, // d2 = oh
+ utils::IteratorType::parallel, // d3 = ow
+ utils::IteratorType::reduction, // d4 = ic_outer
+ utils::IteratorType::reduction, // d5 = fh
+ utils::IteratorType::reduction, // d6 = fw
+ utils::IteratorType::parallel, // d7 = oc_inner (k0)
+ utils::IteratorType::reduction, // d8 = ic_inner (c0)
+ };
+ return space;
+}
+
} // 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 76f816a..b9c509d 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.h
+++ b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.h
@@ -11,6 +11,8 @@
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenTypes.h"
#include "iree/compiler/Dialect/Encoding/IR/EncodingOps.h"
#include "llvm/Support/raw_ostream.h"
+#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
+#include "mlir/IR/AffineMap.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/MLIRContext.h"
#include "mlir/Support/LLVM.h"
@@ -124,6 +126,45 @@
FailureOr<MaterializeEncodingInfo>
getEncodingInfoForMatmul(Encoding::EncodingAttr encoding, TileMxNxK tileMxNxK);
+/// Tile sizes for convolution encoding.
+struct TileOCxIC {
+ int64_t OC = 1; // k0: output channel tile
+ int64_t IC = 1; // c0: input channel tile
+};
+
+/// Returns MaterializeEncodingInfo for a convolution operand.
+FailureOr<MaterializeEncodingInfo>
+getEncodingInfoForConv(Encoding::EncodingAttr encoding, TileOCxIC tile);
+
+//===---------------------------------------------------------------------===//
+// Data-tiled convolution utilities.
+//===---------------------------------------------------------------------===//
+
+/// Iteration space of the data-tiled convolution generic. This is the
+/// canonical definition: the materialization builder and any code that
+/// inspects the resulting op both derive from it.
+///
+/// Loop dims: (n, oc_outer, oh, ow, ic_outer, fh, fw, oc_inner, ic_inner)
+/// d0 d1 d2 d3 d4 d5 d6 d7 d8
+/// input: -> (d0, d4, d2*sH + d5*dH, d3*sW + d6*dW, d8)
+/// filter: -> (d1, d4, d5, d6, d8, d7)
+/// output: -> (d0, d1, d2, d3, d7)
+///
+/// corresponding to packed operand layouts:
+/// input: [N, IC/c0, H, W, c0]
+/// filter: [OC/k0, IC/c0, FH, FW, c0, k0]
+/// output: [N, OC/k0, OH, OW, k0]
+struct DataTiledConvIterationSpace {
+ SmallVector<AffineMap> indexingMaps;
+ SmallVector<utils::IteratorType> iteratorTypes;
+};
+
+/// Returns the indexing maps and iterator types for a data-tiled 2D
+/// convolution with the given strides and dilations.
+DataTiledConvIterationSpace
+getDataTiledConvIterationSpace(MLIRContext *ctx, ArrayRef<int64_t> strides,
+ ArrayRef<int64_t> dilations);
+
} // namespace mlir::iree_compiler::IREE::Codegen
#endif // IREE_COMPILER_CODEGEN_DIALECT_CODEGEN_UTILS_H_
diff --git a/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CPUEncodingExternalModels.cpp b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CPUEncodingExternalModels.cpp
index e3d2324..7260db7 100644
--- a/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CPUEncodingExternalModels.cpp
+++ b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CPUEncodingExternalModels.cpp
@@ -64,6 +64,7 @@
using IREE::Codegen::MaterializeEncodingInfo;
using IREE::Codegen::TileMxNxK;
+using IREE::Codegen::TileOCxIC;
namespace {
@@ -964,6 +965,138 @@
return result;
}
+/// Builds 9D linalg.generic implementing data-tiled convolution
+/// over already-packed, rank-5/6 operands.
+static linalg::GenericOp
+buildDataTiledConvOp(OpBuilder &builder, Location loc, Value packedInput,
+ Value packedFilter, Value packedOutput,
+ ArrayRef<int64_t> strides, ArrayRef<int64_t> dilations) {
+ auto outputType = cast<RankedTensorType>(packedOutput.getType());
+ IREE::Codegen::DataTiledConvIterationSpace space =
+ IREE::Codegen::getDataTiledConvIterationSpace(builder.getContext(),
+ strides, dilations);
+ return linalg::GenericOp::create(
+ builder, loc, outputType,
+ /*inputs=*/ValueRange{packedInput, packedFilter},
+ /*outputs=*/ValueRange{packedOutput}, space.indexingMaps,
+ space.iteratorTypes, [](OpBuilder &b, Location l, ValueRange args) {
+ Value mul = arith::MulFOp::create(b, l, args[0], args[1]);
+ Value add = arith::AddFOp::create(b, l, mul, args[2]);
+ linalg::YieldOp::create(b, l, add);
+ });
+}
+
+/// Lowers a data-tiled convolution whose operands carry conv encodings.
+///
+/// Builds a 9-D linalg.generic over the packed operands:
+/// input: [N, IC/c0, H, W, c0]
+/// filter: [OC/k0, IC/c0, FH, FW, c0, k0]
+/// output: [N, OC/k0, OH, OW, k0]
+Operation *lowerConvolutionOpWithEncoding(
+ OpBuilder &builder, linalg::LinalgOp linalgOp, ValueRange operands,
+ IREE::Encoding::LayoutMaterializerAttr layoutAttr) {
+ if (!linalgOp.hasPureTensorSemantics()) {
+ return nullptr;
+ }
+
+ auto inputs = linalgOp.getDpsInputOperands();
+ auto outputs = linalgOp.getDpsInits();
+
+ auto encodedInputType = cast<RankedTensorType>(inputs[0]->get().getType());
+ auto encodedFilterType = cast<RankedTensorType>(inputs[1]->get().getType());
+ auto encodedOutputType = cast<RankedTensorType>(outputs[0].getType());
+ auto inEncoding = IREE::Encoding::getEncodingAttr(encodedInputType);
+ auto filterEncoding = IREE::Encoding::getEncodingAttr(encodedFilterType);
+ auto outEncoding = IREE::Encoding::getEncodingAttr(encodedOutputType);
+ if (!inEncoding || !filterEncoding || !outEncoding) {
+ return nullptr;
+ }
+
+ if (inEncoding.getOperandIndex().getValue() != IREE::Encoding::CONV_IN ||
+ filterEncoding.getOperandIndex().getValue() !=
+ IREE::Encoding::CONV_FILTER ||
+ outEncoding.getOperandIndex().getValue() != IREE::Encoding::CONV_OUT) {
+ return nullptr;
+ }
+
+ MaterializeEncodingInfo encodingInfo = {};
+ if (auto packedLayoutAttr =
+ dyn_cast<IREE::Codegen::PackedLayoutMaterializerAttr>(layoutAttr)) {
+ encodingInfo = packedLayoutAttr.getEncodingInfo(
+ cast<RankedTensorType>(linalgOp->getResultTypes()[0]));
+ }
+
+ if (isIdentityLayout(encodingInfo)) {
+ return dropEncodingAndCloneOp(builder, linalgOp,
+ operands.take_front(inputs.size()),
+ operands.drop_front(inputs.size()));
+ }
+
+ Value packedInput = operands[0];
+ Value packedFilter = operands[1];
+ Value packedOutput = operands[2];
+
+ Location loc = linalgOp.getLoc();
+
+ auto packedInputType = cast<RankedTensorType>(packedInput.getType());
+ auto packedOutputType = cast<RankedTensorType>(packedOutput.getType());
+
+ // Expected packed ranks for a 2D conv: input/output rank-5 ([N, IC/c0, H, W,
+ // c0] and [N, OC/k0, OH, OW, k0]), or rank-4 if the N=1 batch dim has been
+ // folded away upstream, and filter rank-6 ([OC/k0, IC/c0, FH, FW, c0, k0]).
+ const int64_t inRank = packedInputType.getRank();
+ const int64_t outRank = packedOutputType.getRank();
+ const int64_t filterRank =
+ cast<RankedTensorType>(packedFilter.getType()).getRank();
+ if ((inRank != 4 && inRank != 5) || inRank != outRank || filterRank != 6) {
+ return nullptr;
+ }
+
+ // WORKAROUND: When N=1, upstream folding may remove the batch dimension.
+ // Restore it so the remainder of this lowering can assume rank-5 packed
+ // tensors. If batch folding before materialization is eventually prevented,
+ // this expand/collapse workaround can be dropped.
+ auto prependUnitBatchDim = [&](Value src) -> Value {
+ auto srcType = cast<RankedTensorType>(src.getType());
+
+ SmallVector<int64_t, 5> expandedShape = {1};
+ expandedShape.append(srcType.getShape().begin(), srcType.getShape().end());
+
+ SmallVector<ReassociationIndices> ri = {{0, 1}, {2}, {3}, {4}};
+
+ return tensor::ExpandShapeOp::create(
+ builder, loc,
+ RankedTensorType::get(expandedShape, srcType.getElementType()), src,
+ ri);
+ };
+
+ bool batchDimFolded = (inRank == 4);
+ if (batchDimFolded) {
+ // [IC/c0, H, W, c0] -> [1, IC/c0, H, W, c0]
+ packedInput = prependUnitBatchDim(packedInput);
+ // [OC/k0, OH, OW, k0] -> [1, OC/k0, OH, OW, k0]
+ packedOutput = prependUnitBatchDim(packedOutput);
+ }
+
+ auto cDims = linalg::inferConvolutionDims(linalgOp);
+ if (failed(cDims) || cDims->strides.size() != 2 ||
+ cDims->dilations.size() != 2) {
+ return nullptr;
+ }
+
+ Operation *convOp =
+ buildDataTiledConvOp(builder, loc, packedInput, packedFilter,
+ packedOutput, cDims->strides, cDims->dilations);
+
+ // [1, OC/k0, OH, OW, k0] -> [OC/k0, OH, OW, k0]
+ if (batchDimFolded) {
+ SmallVector<ReassociationIndices> ri = {{0, 1}, {2}, {3}, {4}};
+ return tensor::CollapseShapeOp::create(builder, loc, packedOutputType,
+ convOp->getResult(0), ri);
+ }
+ return convOp;
+}
+
//===----------------------------------------------------------------------===//
// Interface methods implementation for iree_cpu.cpu_encoding_resolver.
//===----------------------------------------------------------------------===//
@@ -1345,6 +1478,57 @@
return {};
}
+// Enumerate (OC, IC) conv inner tile sizes to choose from on arm64.
+static SmallVector<TileOCxIC> enumerateConvTileArm64(TypeRange elementTypes,
+ DictionaryAttr config) {
+ assert(elementTypes.size() == 3 && "expected input, filter, output types");
+ Type inputElem = elementTypes[0];
+ Type filterElem = elementTypes[1];
+ Type outputElem = elementTypes[2];
+ // Only fp32 is currently supported; the tile shape and
+ // the ukernel are specific to 32-bit float lanes.
+ if (inputElem.isF32() && filterElem.isF32() && outputElem.isF32()) {
+ return {TileOCxIC{8, 8}}; // Aim to use FMLA (neon).
+ }
+
+ // Fallback - no architecture-optimized tile size for this case.
+ return {};
+}
+
+// Enumerate (OC, IC) conv inner tile sizes to choose from on x86-64.
+static SmallVector<TileOCxIC> enumerateConvTileX86_64(TypeRange elementTypes,
+ DictionaryAttr config) {
+ assert(elementTypes.size() == 3 && "expected input, filter, output types");
+ Type inputElem = elementTypes[0];
+ Type filterElem = elementTypes[1];
+ Type outputElem = elementTypes[2];
+ // Only fp32 is currently supported; the tile shape and
+ // the ukernel are specific to 32-bit float lanes.
+ if (inputElem.isF32() && filterElem.isF32() && outputElem.isF32()) {
+ if (hasFeature(config, "+avx512f")) {
+ return {TileOCxIC{16, 16}}; // Aim to use VFMADD* (zmm).
+ }
+ }
+
+ // Fallback - no architecture-optimized tile size for this case.
+ return {};
+}
+
+static SmallVector<TileOCxIC>
+enumerateCPUConvTiles(IREE::Encoding::EncodingAttr encoding,
+ DictionaryAttr config) {
+ SmallVector<Type> elementTypes = encoding.getElementTypesArray();
+ if (isAArch64(config)) {
+ return enumerateConvTileArm64(elementTypes, config);
+ }
+ if (isX86_64(config)) {
+ return enumerateConvTileX86_64(elementTypes, config);
+ }
+
+ // Fallback - no architecture-optimized tile size for this case.
+ return {};
+}
+
struct CPUEncodingPackedLayoutMaterializerAttr
: PackedLayoutMaterializerAttrExternalModelBase<
CPUEncodingPackedLayoutMaterializerAttr, CPUEncodingResolverAttr> {
@@ -1365,6 +1549,30 @@
return info;
}
+ // Handle convolution encodings.
+ if (encoding.getOpType().getValue() ==
+ IREE::Encoding::EncodingOpType::conv) {
+ FailureOr<linalg::ConvolutionDimensions> cDims =
+ linalg::inferConvolutionDims(encoding.getRootMaps());
+ if (failed(cDims) || !cDims->depth.empty() ||
+ cDims->outputImage.size() != 2 ||
+ llvm::any_of(cDims->dilations, [](int64_t d) { return d != 1; })) {
+ return info; // identity layout
+ }
+ SmallVector<TileOCxIC> tiles =
+ enumerateCPUConvTiles(encoding, layoutAttr.getConfiguration());
+ if (tiles.empty()) {
+ return info;
+ }
+ TileOCxIC tile = tiles[0];
+ FailureOr<MaterializeEncodingInfo> maybeInfo =
+ IREE::Codegen::getEncodingInfoForConv(encoding, tile);
+ if (failed(maybeInfo)) {
+ return info;
+ }
+ return maybeInfo.value();
+ }
+
// We only know about contractions with {Batch, M, N, K} <= 1 at the moment.
auto cDims = getEncodingContractionDims(encoding);
if (failed(cDims) || cDims->batch.size() > 1 || cDims->m.size() > 1 ||
@@ -1496,6 +1704,11 @@
convertedOperands.take_front(numInputs),
convertedOperands.drop_front(numInputs));
}
+ if (linalg::isaConvolutionOpInterface(linalgOp)) {
+ return lowerConvolutionOpWithEncoding(
+ b, linalgOp, convertedOperands,
+ cast<IREE::Encoding::LayoutMaterializerAttr>(layoutAttr));
+ }
if (linalg::isaContractionOpInterface(linalgOp)) {
DictionaryAttr config = layoutAttr.getConfiguration();
if (getEnableInnerTiledFromConfig(config)) {
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/specialize_encodings.mlir b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/specialize_encodings.mlir
index f993c71..36b0e06 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/specialize_encodings.mlir
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/specialize_encodings.mlir
@@ -40,6 +40,39 @@
// -----
+// Encoding layout specialization for an nhwc_hwcf convolution.
+
+#conv_in = affine_map<(n, oh, ow, oc, fh, fw, ic) -> (n, oh + fh, ow + fw, ic)>
+#conv_filter = affine_map<(n, oh, ow, oc, fh, fw, ic) -> (fh, fw, ic, oc)>
+#conv_out = affine_map<(n, oh, ow, oc, fh, fw, ic) -> (n, oh, ow, oc)>
+#executable_target_x86_64 = #hal.executable.target<"llvm-cpu", "xyz", {iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>, target_triple="x86_64-xyz-xyz", cpu_features="+avx512f"}>
+#executable_target_aarch64 = #hal.executable.target<"llvm-cpu", "xyz", {iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>, target_triple="aarch64-xyz-xyz", cpu_features="+neon"}>
+#device_target_local_0_ = #hal.device.target<"local", {ordinal = 0 : index}, [#executable_target_x86_64]> : !hal.device
+#device_target_local_1_ = #hal.device.target<"local", {ordinal = 1 : index}, [#executable_target_aarch64]> : !hal.device
+#conv_encoding_input = #iree_encoding.encoding<operand_index = 0 : index, op_type = conv, element_types = [f32, f32, f32], user_indexing_maps = [#conv_in, #conv_filter, #conv_out], iteration_sizes = [1, 14, 14, 8, 3, 3, 4]>
+#conv_encoding_filter = #iree_encoding.encoding<operand_index = 1 : index, op_type = conv, element_types = [f32, f32, f32], user_indexing_maps = [#conv_in, #conv_filter, #conv_out], iteration_sizes = [1, 14, 14, 8, 3, 3, 4]>
+
+util.global private @device_a = #device_target_local_0_
+util.global private @device_b = #device_target_local_1_
+util.func public @conv_tensor_sizeof() -> (index, index, index, index) {
+ %size0 = stream.tensor.sizeof on(#hal.device.affinity<@device_a>) tensor<1x16x16x4xf32, #conv_encoding_input> : index
+ %size1 = stream.tensor.sizeof on(#hal.device.affinity<@device_a>) tensor<3x3x4x8xf32, #conv_encoding_filter> : index
+ %size2 = stream.tensor.sizeof on(#hal.device.affinity<@device_b>) tensor<1x16x16x4xf32, #conv_encoding_input> : index
+ %size3 = stream.tensor.sizeof on(#hal.device.affinity<@device_b>) tensor<3x3x4x8xf32, #conv_encoding_filter> : index
+ util.return %size0, %size1, %size2, %size3 : index, index, index, index
+}
+// CHECK-DAG: #[[$IN_X86:.+]] = #iree_encoding.layout<[#iree_cpu.cpu_encoding_resolver{{.+}}encoding_info = {innerDimsPos = [3], innerTileSizes = [16], outerDimsPerm = [0, 3, 1, 2]}
+// CHECK-DAG: #[[$FLT_X86:.+]] = #iree_encoding.layout<[#iree_cpu.cpu_encoding_resolver{{.+}}encoding_info = {innerDimsPos = [2, 3], innerTileSizes = [16, 16], outerDimsPerm = [3, 2, 0, 1]}
+// CHECK-DAG: #[[$IN_ARM:.+]] = #iree_encoding.layout<[#iree_cpu.cpu_encoding_resolver{{.+}}encoding_info = {innerDimsPos = [3], innerTileSizes = [8], outerDimsPerm = [0, 3, 1, 2]}
+// CHECK-DAG: #[[$FLT_ARM:.+]] = #iree_encoding.layout<[#iree_cpu.cpu_encoding_resolver{{.+}}encoding_info = {innerDimsPos = [2, 3], innerTileSizes = [8, 8], outerDimsPerm = [3, 2, 0, 1]}
+// CHECK-LABEL: util.func public @conv_tensor_sizeof
+// CHECK: stream.tensor.sizeof {{.+}} tensor<1x16x16x4xf32, #[[$IN_X86]]>
+// CHECK: stream.tensor.sizeof {{.+}} tensor<3x3x4x8xf32, #[[$FLT_X86]]>
+// CHECK: stream.tensor.sizeof {{.+}} tensor<1x16x16x4xf32, #[[$IN_ARM]]>
+// CHECK: stream.tensor.sizeof {{.+}} tensor<3x3x4x8xf32, #[[$FLT_ARM]]>
+
+// -----
+
//------------------------------------------------------------------------------
// #iree_gpu.gpu_encoding_resolver specialization tests.
// These get serialized to the layout attributes.
