[Runtime] Add initial data-tiled conv ukernel (generic & avx512) (#24729)
Adds a data-tiled convolution microkernel with a generic implementation
and an avx512 optimized tile kernel with inner tile sizes (OW_TILE=16,
K0=16, C0=16).
Layouts (c0 = inner IC block, k0 = inner OC block):
input [N, IC/c0, H, W, c0]
filter [OC/k0, IC/c0, FH, FW, c0, k0]
output [N, OC/k0, OH, OW, k0]
The ABI uses explicit outer-dimension strides. These strides refer to
the full parent tensor layout rather than the extracted slab, ensuring
correct addressing when kernels operate on subviews of the tensor. Only
the inner W/c0/k0 dims are contiguous.
Scope: f32 only; dilation=1 (dilated convs are not matched by the
dispatch); requires the packed NCHWc layout, with the arch path selected
only for k0=c0=16. Additional types and tile shapes are future work.
---------
Signed-off-by: Pooja Hemashekar <hemashekar@roofline.ai>diff --git a/runtime/src/iree/builtins/ukernel/BUILD.bazel b/runtime/src/iree/builtins/ukernel/BUILD.bazel
index 22448d4..2a9efb4 100644
--- a/runtime/src/iree/builtins/ukernel/BUILD.bazel
+++ b/runtime/src/iree/builtins/ukernel/BUILD.bazel
@@ -21,6 +21,8 @@
internal_headers = [
"common.h",
+ "conv_nchwc.h",
+ "conv_nchwc_internal.h",
"exported_bits.h",
"mmt4d.h",
"mmt4d_internal.h",
@@ -58,6 +60,8 @@
iree_runtime_cc_library(
name = "ukernel",
srcs = [
+ "conv_nchwc.c",
+ "conv_nchwc_tile_generic.c",
"mmt4d.c",
"mmt4d_tile_generic.c",
"pack.c",
@@ -112,6 +116,8 @@
[iree_bitcode_library(
name = "ukernel_bitcode_generic_%s" % arch,
srcs = [
+ "conv_nchwc.c",
+ "conv_nchwc_tile_generic.c",
"mmt4d.c",
"mmt4d_tile_generic.c",
"pack.c",
diff --git a/runtime/src/iree/builtins/ukernel/CMakeLists.txt b/runtime/src/iree/builtins/ukernel/CMakeLists.txt
index 1689c49..ba8429a 100644
--- a/runtime/src/iree/builtins/ukernel/CMakeLists.txt
+++ b/runtime/src/iree/builtins/ukernel/CMakeLists.txt
@@ -33,6 +33,8 @@
COMMAND ${CMAKE_COMMAND} -E touch internal_headers_filegroup.stamp
DEPENDS
"common.h"
+ "conv_nchwc.h"
+ "conv_nchwc_internal.h"
"exported_bits.h"
"mmt4d.h"
"mmt4d_internal.h"
@@ -53,6 +55,8 @@
internal_headers
HDRS
"common.h"
+ "conv_nchwc.h"
+ "conv_nchwc_internal.h"
"exported_bits.h"
"mmt4d.h"
"mmt4d_internal.h"
@@ -72,6 +76,8 @@
fallback
HDRS
"common.h"
+ "conv_nchwc.h"
+ "conv_nchwc_internal.h"
"exported_bits.h"
"mmt4d.h"
"mmt4d_internal.h"
@@ -95,6 +101,10 @@
"api.h"
SRCS
"common.h"
+ "conv_nchwc.c"
+ "conv_nchwc.h"
+ "conv_nchwc_internal.h"
+ "conv_nchwc_tile_generic.c"
"exported_bits.h"
"mmt4d.c"
"mmt4d.h"
@@ -129,6 +139,8 @@
"${PROJECT_BINARY_DIR}/runtime/src/iree/schemas/cpu_data_headers_filegroup.stamp"
"internal_headers_filegroup.stamp"
SRCS
+ "conv_nchwc.c"
+ "conv_nchwc_tile_generic.c"
"mmt4d.c"
"mmt4d_tile_generic.c"
"pack.c"
@@ -146,6 +158,8 @@
"${PROJECT_BINARY_DIR}/runtime/src/iree/schemas/cpu_data_headers_filegroup.stamp"
"internal_headers_filegroup.stamp"
SRCS
+ "conv_nchwc.c"
+ "conv_nchwc_tile_generic.c"
"mmt4d.c"
"mmt4d_tile_generic.c"
"pack.c"
@@ -163,6 +177,8 @@
"${PROJECT_BINARY_DIR}/runtime/src/iree/schemas/cpu_data_headers_filegroup.stamp"
"internal_headers_filegroup.stamp"
SRCS
+ "conv_nchwc.c"
+ "conv_nchwc_tile_generic.c"
"fallback.c"
"mmt4d.c"
"mmt4d_tile_generic.c"
@@ -181,6 +197,8 @@
"${PROJECT_BINARY_DIR}/runtime/src/iree/schemas/cpu_data_headers_filegroup.stamp"
"internal_headers_filegroup.stamp"
SRCS
+ "conv_nchwc.c"
+ "conv_nchwc_tile_generic.c"
"mmt4d.c"
"mmt4d_tile_generic.c"
"pack.c"
@@ -198,6 +216,8 @@
"${PROJECT_BINARY_DIR}/runtime/src/iree/schemas/cpu_data_headers_filegroup.stamp"
"internal_headers_filegroup.stamp"
SRCS
+ "conv_nchwc.c"
+ "conv_nchwc_tile_generic.c"
"fallback.c"
"mmt4d.c"
"mmt4d_tile_generic.c"
diff --git a/runtime/src/iree/builtins/ukernel/api.h b/runtime/src/iree/builtins/ukernel/api.h
index a6584b3..b557d00 100644
--- a/runtime/src/iree/builtins/ukernel/api.h
+++ b/runtime/src/iree/builtins/ukernel/api.h
@@ -7,6 +7,7 @@
#ifndef IREE_BUILTINS_UKERNEL_API_H_
#define IREE_BUILTINS_UKERNEL_API_H_
+#include "iree/builtins/ukernel/conv_nchwc.h"
#include "iree/builtins/ukernel/mmt4d.h"
#include "iree/builtins/ukernel/pack.h"
#include "iree/builtins/ukernel/query_tile_sizes.h"
diff --git a/runtime/src/iree/builtins/ukernel/arch/arm_64/BUILD.bazel b/runtime/src/iree/builtins/ukernel/arch/arm_64/BUILD.bazel
index 182b9b3..2492d0d 100644
--- a/runtime/src/iree/builtins/ukernel/arch/arm_64/BUILD.bazel
+++ b/runtime/src/iree/builtins/ukernel/arch/arm_64/BUILD.bazel
@@ -39,6 +39,7 @@
iree_bitcode_library(
name = "ukernel_bitcode_arch_arm_64_entry_points",
srcs = [
+ "conv_nchwc_arm_64_entry_point.c",
"mmt4d_arm_64_entry_point.c",
"pack_arm_64_entry_point.c",
"unpack_arm_64_entry_point.c",
diff --git a/runtime/src/iree/builtins/ukernel/arch/arm_64/CMakeLists.txt b/runtime/src/iree/builtins/ukernel/arch/arm_64/CMakeLists.txt
index c10a595..72d2f1c 100644
--- a/runtime/src/iree/builtins/ukernel/arch/arm_64/CMakeLists.txt
+++ b/runtime/src/iree/builtins/ukernel/arch/arm_64/CMakeLists.txt
@@ -27,6 +27,7 @@
"pack_arm_64_internal.h"
"unpack_arm_64_internal.h"
SRCS
+ "conv_nchwc_arm_64_entry_point.c"
"mmt4d_arm_64_entry_point.c"
"pack_arm_64_entry_point.c"
"unpack_arm_64_entry_point.c"
@@ -288,6 +289,7 @@
NAME
arm_64
SRCS
+ "conv_nchwc_arm_64_entry_point.c"
"mmt4d_arm_64_entry_point.c"
"mmt4d_arm_64_base.c"
"pack_arm_64_entry_point.c"
diff --git a/runtime/src/iree/builtins/ukernel/arch/arm_64/conv_nchwc_arm_64_entry_point.c b/runtime/src/iree/builtins/ukernel/arch/arm_64/conv_nchwc_arm_64_entry_point.c
new file mode 100644
index 0000000..dc9ef7d
--- /dev/null
+++ b/runtime/src/iree/builtins/ukernel/arch/arm_64/conv_nchwc_arm_64_entry_point.c
@@ -0,0 +1,17 @@
+// Copyright 2026 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+// No arm_64-specific conv_nchwc tile function yet. Falls back to generic.
+// TODO: implement arm_64 tile functions analogous to the +dotprod / +i8mm /
+// +bf16 dispatch in mmt4d_arm_64_entry_point.c.
+#include "iree/builtins/ukernel/conv_nchwc_internal.h"
+
+iree_uk_conv_nchwc_tile_selection_t iree_uk_conv_nchwc_select_tile_func_arch(
+ const iree_uk_conv_nchwc_params_t* params) {
+ (void)params;
+ iree_uk_conv_nchwc_tile_selection_t selection = {0};
+ return selection;
+}
diff --git a/runtime/src/iree/builtins/ukernel/arch/riscv_64/BUILD.bazel b/runtime/src/iree/builtins/ukernel/arch/riscv_64/BUILD.bazel
index b41e2f4..02dae9e 100644
--- a/runtime/src/iree/builtins/ukernel/arch/riscv_64/BUILD.bazel
+++ b/runtime/src/iree/builtins/ukernel/arch/riscv_64/BUILD.bazel
@@ -39,6 +39,7 @@
iree_bitcode_library(
name = "ukernel_bitcode_arch_riscv_64_entry_points",
srcs = [
+ "conv_nchwc_riscv_64_entry_point.c",
"mmt4d_riscv_64_entry_point.c",
"pack_riscv_64_entry_point.c",
"unpack_riscv_64_entry_point.c",
diff --git a/runtime/src/iree/builtins/ukernel/arch/riscv_64/CMakeLists.txt b/runtime/src/iree/builtins/ukernel/arch/riscv_64/CMakeLists.txt
index 7cd59ef..f155c06 100644
--- a/runtime/src/iree/builtins/ukernel/arch/riscv_64/CMakeLists.txt
+++ b/runtime/src/iree/builtins/ukernel/arch/riscv_64/CMakeLists.txt
@@ -27,6 +27,7 @@
"pack_riscv_64_internal.h"
"unpack_riscv_64_internal.h"
SRCS
+ "conv_nchwc_riscv_64_entry_point.c"
"mmt4d_riscv_64_entry_point.c"
"pack_riscv_64_entry_point.c"
"unpack_riscv_64_entry_point.c"
@@ -259,6 +260,7 @@
NAME
riscv_64
SRCS
+ "conv_nchwc_riscv_64_entry_point.c"
"mmt4d_riscv_64_entry_point.c"
"pack_riscv_64_entry_point.c"
"unpack_riscv_64_entry_point.c"
diff --git a/runtime/src/iree/builtins/ukernel/arch/riscv_64/conv_nchwc_riscv_64_entry_point.c b/runtime/src/iree/builtins/ukernel/arch/riscv_64/conv_nchwc_riscv_64_entry_point.c
new file mode 100644
index 0000000..6860e6e
--- /dev/null
+++ b/runtime/src/iree/builtins/ukernel/arch/riscv_64/conv_nchwc_riscv_64_entry_point.c
@@ -0,0 +1,17 @@
+// Copyright 2026 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+// No riscv_64-specific conv_nchwc tile function yet. Falls back to generic.
+// TODO: implement riscv_64 tile functions analogous to the +v / +zvfh /
+// +zvfhmin dispatch in mmt4d_riscv_64_entry_point.c.
+#include "iree/builtins/ukernel/conv_nchwc_internal.h"
+
+iree_uk_conv_nchwc_tile_selection_t iree_uk_conv_nchwc_select_tile_func_arch(
+ const iree_uk_conv_nchwc_params_t* params) {
+ (void)params;
+ iree_uk_conv_nchwc_tile_selection_t selection = {0};
+ return selection;
+}
diff --git a/runtime/src/iree/builtins/ukernel/arch/x86_64/BUILD.bazel b/runtime/src/iree/builtins/ukernel/arch/x86_64/BUILD.bazel
index f487707..28f8a6f 100644
--- a/runtime/src/iree/builtins/ukernel/arch/x86_64/BUILD.bazel
+++ b/runtime/src/iree/builtins/ukernel/arch/x86_64/BUILD.bazel
@@ -28,6 +28,8 @@
# All headers transitively included by code in this directory. Bazel-only.
UKERNEL_X86_64_INTERNAL_HEADERS = [
"common_x86_64.h",
+ "conv_nchwc_x86_64_internal.h",
+ "conv_nchwc_x86_64_tiles.inl",
"mmt4d_x86_64_internal.h",
"mmt4d_x86_64_tiles.inl",
"pack_x86_64_internal.h",
@@ -39,6 +41,7 @@
iree_bitcode_library(
name = "ukernel_bitcode_arch_x86_64_entry_points",
srcs = [
+ "conv_nchwc_x86_64_entry_point.c",
"mmt4d_x86_64_entry_point.c",
"pack_x86_64_entry_point.c",
"unpack_x86_64_entry_point.c",
@@ -77,6 +80,7 @@
iree_bitcode_library(
name = "ukernel_bitcode_arch_x86_64_avx512_base",
srcs = [
+ "conv_nchwc_x86_64_avx512_base.c",
"mmt4d_x86_64_avx512_base.c",
"pack_x86_64_avx512_base.c",
"unpack_x86_64_avx512_base.c",
diff --git a/runtime/src/iree/builtins/ukernel/arch/x86_64/CMakeLists.txt b/runtime/src/iree/builtins/ukernel/arch/x86_64/CMakeLists.txt
index 63d1376..ec84021 100644
--- a/runtime/src/iree/builtins/ukernel/arch/x86_64/CMakeLists.txt
+++ b/runtime/src/iree/builtins/ukernel/arch/x86_64/CMakeLists.txt
@@ -22,11 +22,14 @@
"${PROJECT_BINARY_DIR}/runtime/src/iree/builtins/ukernel/internal_headers_filegroup.stamp"
"${PROJECT_BINARY_DIR}/runtime/src/iree/schemas/cpu_data_headers_filegroup.stamp"
"common_x86_64.h"
+ "conv_nchwc_x86_64_internal.h"
+ "conv_nchwc_x86_64_tiles.inl"
"mmt4d_x86_64_internal.h"
"mmt4d_x86_64_tiles.inl"
"pack_x86_64_internal.h"
"unpack_x86_64_internal.h"
SRCS
+ "conv_nchwc_x86_64_entry_point.c"
"mmt4d_x86_64_entry_point.c"
"pack_x86_64_entry_point.c"
"unpack_x86_64_entry_point.c"
@@ -41,6 +44,8 @@
"${PROJECT_BINARY_DIR}/runtime/src/iree/builtins/ukernel/internal_headers_filegroup.stamp"
"${PROJECT_BINARY_DIR}/runtime/src/iree/schemas/cpu_data_headers_filegroup.stamp"
"common_x86_64.h"
+ "conv_nchwc_x86_64_internal.h"
+ "conv_nchwc_x86_64_tiles.inl"
"mmt4d_x86_64_internal.h"
"mmt4d_x86_64_tiles.inl"
"pack_x86_64_internal.h"
@@ -65,11 +70,14 @@
"${PROJECT_BINARY_DIR}/runtime/src/iree/builtins/ukernel/internal_headers_filegroup.stamp"
"${PROJECT_BINARY_DIR}/runtime/src/iree/schemas/cpu_data_headers_filegroup.stamp"
"common_x86_64.h"
+ "conv_nchwc_x86_64_internal.h"
+ "conv_nchwc_x86_64_tiles.inl"
"mmt4d_x86_64_internal.h"
"mmt4d_x86_64_tiles.inl"
"pack_x86_64_internal.h"
"unpack_x86_64_internal.h"
SRCS
+ "conv_nchwc_x86_64_avx512_base.c"
"mmt4d_x86_64_avx512_base.c"
"pack_x86_64_avx512_base.c"
"unpack_x86_64_avx512_base.c"
@@ -94,6 +102,8 @@
"${PROJECT_BINARY_DIR}/runtime/src/iree/builtins/ukernel/internal_headers_filegroup.stamp"
"${PROJECT_BINARY_DIR}/runtime/src/iree/schemas/cpu_data_headers_filegroup.stamp"
"common_x86_64.h"
+ "conv_nchwc_x86_64_internal.h"
+ "conv_nchwc_x86_64_tiles.inl"
"mmt4d_x86_64_internal.h"
"mmt4d_x86_64_tiles.inl"
"pack_x86_64_internal.h"
@@ -122,6 +132,8 @@
"${PROJECT_BINARY_DIR}/runtime/src/iree/builtins/ukernel/internal_headers_filegroup.stamp"
"${PROJECT_BINARY_DIR}/runtime/src/iree/schemas/cpu_data_headers_filegroup.stamp"
"common_x86_64.h"
+ "conv_nchwc_x86_64_internal.h"
+ "conv_nchwc_x86_64_tiles.inl"
"mmt4d_x86_64_internal.h"
"mmt4d_x86_64_tiles.inl"
"pack_x86_64_internal.h"
@@ -367,6 +379,7 @@
NAME
x86_64_avx512_base
SRCS
+ "conv_nchwc_x86_64_avx512_base.c"
"mmt4d_x86_64_avx512_base.c"
"pack_x86_64_avx512_base.c"
"unpack_x86_64_avx512_base.c"
@@ -410,6 +423,7 @@
NAME
x86_64
SRCS
+ "conv_nchwc_x86_64_entry_point.c"
"mmt4d_x86_64_entry_point.c"
"pack_x86_64_entry_point.c"
"query_tile_sizes_x86_64_entry_point.c"
diff --git a/runtime/src/iree/builtins/ukernel/arch/x86_64/conv_nchwc_x86_64_avx512_base.c b/runtime/src/iree/builtins/ukernel/arch/x86_64/conv_nchwc_x86_64_avx512_base.c
new file mode 100644
index 0000000..7b5f61e
--- /dev/null
+++ b/runtime/src/iree/builtins/ukernel/arch/x86_64/conv_nchwc_x86_64_avx512_base.c
@@ -0,0 +1,118 @@
+// Copyright 2026 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/builtins/ukernel/arch/x86_64/common_x86_64.h"
+#include "iree/builtins/ukernel/arch/x86_64/conv_nchwc_x86_64_internal.h"
+
+// AVX-512 tile kernel for data-tiled NCHWc conv (f32 x f32 -> f32).
+//
+// 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]
+//
+// Computes one output panel of shape [OW_TILE=16, k0=16] per call.
+//
+// Strategy:
+// - 16 ZMM accumulators, one per output-width position, each holding a
+// k0-wide OC block.
+// - For each (ic_outer, fh, fw, ci):
+// load one ZMM containing k0 output-channel values for fixed (ic_outer,
+// fh, fw, ci), broadcast scalar input value at (ow, fw, ci) and FMA into
+// each OW accumulator.
+//
+// Tile shapes: OW_TILE=16, k0=16, c0=16.
+#define OW_TILE 16
+#define K0 16
+#define C0 16
+
+IREE_UK_ATTRIBUTE_ALWAYS_INLINE static inline void
+iree_uk_conv_nchwc_tile_f32f32f32_16x16x16_x86_64_avx512_base_impl(
+ float* IREE_UK_RESTRICT out_ptr, const float* IREE_UK_RESTRICT in_ptr,
+ const float* IREE_UK_RESTRICT f_ptr,
+ const iree_uk_conv_nchwc_params_t* params, iree_uk_index_t ow_count) {
+ const iree_uk_index_t IC_outer = params->IC_outer;
+ const iree_uk_index_t FH = params->FH;
+ const iree_uk_index_t FW = params->FW;
+ const iree_uk_index_t in_pos_stride = (iree_uk_index_t)params->stride_w * C0;
+ const iree_uk_index_t in_stride_h = params->input_stride_h;
+ const iree_uk_index_t in_stride_ic_outer = params->input_stride_ic_outer;
+ const iree_uk_index_t f_stride_fw = (iree_uk_index_t)C0 * K0;
+ const iree_uk_index_t f_stride_fh = params->filter_stride_fh;
+ const iree_uk_index_t f_stride_ic_outer = params->filter_stride_ic_outer;
+
+ __m512 acc[OW_TILE];
+ if (params->flags & IREE_UK_FLAG_CONV_NCHWC_ACCUMULATE) {
+ IREE_UK_UNROLL for (int i = 0; i < OW_TILE; ++i) {
+ if (i < ow_count) {
+ acc[i] = _mm512_loadu_ps(out_ptr + i * K0);
+ } else {
+ acc[i] = _mm512_setzero_ps();
+ }
+ }
+ } else {
+ IREE_UK_UNROLL for (int i = 0; i < OW_TILE; ++i) {
+ acc[i] = _mm512_setzero_ps();
+ }
+ }
+
+ for (iree_uk_index_t ic_o = 0; ic_o < IC_outer; ++ic_o) {
+ const float* in_ic = in_ptr + ic_o * in_stride_ic_outer;
+ const float* f_ic = f_ptr + ic_o * f_stride_ic_outer;
+
+ for (iree_uk_index_t fh = 0; fh < FH; ++fh) {
+ const float* in_h = in_ic + fh * in_stride_h;
+ const float* f_h = f_ic + fh * f_stride_fh;
+
+ for (iree_uk_index_t fw = 0; fw < FW; ++fw) {
+ const float* f_fw = f_h + fw * f_stride_fw;
+ const float* in_fw = in_h + fw * C0;
+ for (int ci = 0; ci < C0; ++ci) {
+ // Load filter[ci, 0:k0].
+ __m512 f_vec = _mm512_loadu_ps(f_fw + ci * K0);
+
+ // Broadcast input[..., ci] for each output position and FMA.
+ IREE_UK_UNROLL for (int i = 0; i < OW_TILE; ++i) {
+ if (i < ow_count) {
+ __m512 in_bc = _mm512_set1_ps(in_fw[i * in_pos_stride + ci]);
+ acc[i] = _mm512_fmadd_ps(in_bc, f_vec, acc[i]);
+ }
+ }
+ }
+ }
+ }
+ }
+
+ IREE_UK_UNROLL for (int i = 0; i < OW_TILE; ++i) {
+ if (i < ow_count) {
+ _mm512_storeu_ps(out_ptr + i * K0, acc[i]);
+ }
+ }
+}
+
+void iree_uk_conv_nchwc_tile_f32f32f32_16x16x16_x86_64_avx512_base(
+ void* IREE_UK_RESTRICT output_panel,
+ const void* IREE_UK_RESTRICT input_panel,
+ const void* IREE_UK_RESTRICT filter_panel,
+ const iree_uk_conv_nchwc_params_t* params) {
+ IREE_UK_ASSERT(params->ow_count > 0 && params->ow_count <= OW_TILE);
+ float* IREE_UK_RESTRICT out_ptr = (float*)output_panel;
+ const float* IREE_UK_RESTRICT in_ptr = (const float*)input_panel;
+ const float* IREE_UK_RESTRICT f_ptr = (const float*)filter_panel;
+ // TODO(phemashekar): adopt the "always pad" rule (ow_count == OW_TILE always)
+ // to delete the tail path entirely.
+ if (params->ow_count == OW_TILE) {
+ iree_uk_conv_nchwc_tile_f32f32f32_16x16x16_x86_64_avx512_base_impl(
+ out_ptr, in_ptr, f_ptr, params, OW_TILE);
+ } else {
+ iree_uk_conv_nchwc_tile_f32f32f32_16x16x16_x86_64_avx512_base_impl(
+ out_ptr, in_ptr, f_ptr, params, params->ow_count);
+ }
+}
+
+#undef OW_TILE
+#undef K0
+#undef C0
diff --git a/runtime/src/iree/builtins/ukernel/arch/x86_64/conv_nchwc_x86_64_entry_point.c b/runtime/src/iree/builtins/ukernel/arch/x86_64/conv_nchwc_x86_64_entry_point.c
new file mode 100644
index 0000000..925a835
--- /dev/null
+++ b/runtime/src/iree/builtins/ukernel/arch/x86_64/conv_nchwc_x86_64_entry_point.c
@@ -0,0 +1,44 @@
+// Copyright 2026 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/builtins/ukernel/arch/x86_64/common_x86_64.h"
+#include "iree/builtins/ukernel/arch/x86_64/conv_nchwc_x86_64_internal.h"
+
+iree_uk_conv_nchwc_tile_selection_t iree_uk_conv_nchwc_select_tile_func_arch(
+ const iree_uk_conv_nchwc_params_t* params) {
+ IREE_UK_ATTRIBUTE_UNUSED iree_uk_conv_nchwc_type_t type =
+ iree_uk_conv_nchwc_type(params->flags);
+ iree_uk_conv_nchwc_tile_selection_t selection = {0};
+
+#define IREE_UK_CONV_NCHWC_TILE_IMPL_x86_64(LHS, RHS, OUT, OW_TILE, K0_TILE, \
+ C0_TILE, SUFFIX) \
+ if (type == iree_uk_conv_nchwc_type_##LHS##RHS##OUT && \
+ params->k0 == (K0_TILE) && params->c0 == (C0_TILE) && \
+ iree_uk_cpu_x86_64##SUFFIX(params->cpu_data)) { \
+ selection.tile_func = \
+ iree_uk_conv_nchwc_tile_##LHS##RHS##OUT##_##OW_TILE##x##K0_TILE##x##C0_TILE##_x86_64##SUFFIX; \
+ selection.ow_tile = (OW_TILE); \
+ }
+
+#ifdef IREE_UK_BUILD_X86_64_AVX512_BASE
+#define IREE_UK_CONV_NCHWC_TILE_x86_64_avx512_base(LHS, RHS, OUT, OW_TILE, \
+ K0_TILE, C0_TILE) \
+ IREE_UK_CONV_NCHWC_TILE_IMPL_x86_64(LHS, RHS, OUT, OW_TILE, K0_TILE, \
+ C0_TILE, _avx512_base)
+#else
+#define IREE_UK_CONV_NCHWC_TILE_x86_64_avx512_base(LHS, RHS, OUT, OW_TILE, \
+ K0_TILE, C0_TILE)
+#endif
+
+#define IREE_UK_CONV_NCHWC_TILE(ARCH, LHS, RHS, OUT, OW_TILE, K0_TILE, \
+ C0_TILE, SUFFIX) \
+ IREE_UK_CONV_NCHWC_TILE_x86_64##SUFFIX(LHS, RHS, OUT, OW_TILE, K0_TILE, \
+ C0_TILE)
+
+#include "iree/builtins/ukernel/arch/x86_64/conv_nchwc_x86_64_tiles.inl"
+
+ return selection;
+}
diff --git a/runtime/src/iree/builtins/ukernel/arch/x86_64/conv_nchwc_x86_64_internal.h b/runtime/src/iree/builtins/ukernel/arch/x86_64/conv_nchwc_x86_64_internal.h
new file mode 100644
index 0000000..c1a29e5
--- /dev/null
+++ b/runtime/src/iree/builtins/ukernel/arch/x86_64/conv_nchwc_x86_64_internal.h
@@ -0,0 +1,24 @@
+// Copyright 2026 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#ifndef IREE_BUILTINS_UKERNEL_ARCH_X86_64_CONV_NCHWC_X86_64_INTERNAL_H_
+#define IREE_BUILTINS_UKERNEL_ARCH_X86_64_CONV_NCHWC_X86_64_INTERNAL_H_
+
+#include "iree/builtins/ukernel/conv_nchwc_internal.h"
+
+#define IREE_UK_CONV_NCHWC_TILE(ARCH, LHS, RHS, OUT, OW_TILE, K0, C0, SUFFIX) \
+ void \
+ iree_uk_conv_nchwc_tile_##LHS##RHS##OUT##_##OW_TILE##x##K0##x##C0##_##ARCH##SUFFIX( \
+ void* IREE_UK_RESTRICT output_panel, \
+ const void* IREE_UK_RESTRICT input_panel, \
+ const void* IREE_UK_RESTRICT filter_panel, \
+ const iree_uk_conv_nchwc_params_t* params);
+
+#include "iree/builtins/ukernel/arch/x86_64/conv_nchwc_x86_64_tiles.inl"
+
+#undef IREE_UK_CONV_NCHWC_TILE
+
+#endif // IREE_BUILTINS_UKERNEL_ARCH_X86_64_CONV_NCHWC_X86_64_INTERNAL_H_
diff --git a/runtime/src/iree/builtins/ukernel/arch/x86_64/conv_nchwc_x86_64_tiles.inl b/runtime/src/iree/builtins/ukernel/arch/x86_64/conv_nchwc_x86_64_tiles.inl
new file mode 100644
index 0000000..e0a47b2
--- /dev/null
+++ b/runtime/src/iree/builtins/ukernel/arch/x86_64/conv_nchwc_x86_64_tiles.inl
@@ -0,0 +1,11 @@
+// Copyright 2026 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+// Params: ARCH, INPUT_T, FILTER_T, OUTPUT_T, OW_TILE, K0, C0, SUFFIX
+//
+// AVX-512 provides a single f32 tile config for now:
+// OW_TILE=16, K0=16, C0=16
+IREE_UK_CONV_NCHWC_TILE(x86_64, f32, f32, f32, 16, 16, 16, _avx512_base)
diff --git a/runtime/src/iree/builtins/ukernel/conv_nchwc.c b/runtime/src/iree/builtins/ukernel/conv_nchwc.c
new file mode 100644
index 0000000..68d101d
--- /dev/null
+++ b/runtime/src/iree/builtins/ukernel/conv_nchwc.c
@@ -0,0 +1,219 @@
+// Copyright 2026 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/builtins/ukernel/conv_nchwc.h"
+
+#include "iree/builtins/ukernel/conv_nchwc_internal.h"
+#include "iree/builtins/ukernel/exported_bits.h"
+
+static void iree_uk_conv_nchwc_validate(
+ const iree_uk_conv_nchwc_params_t* params) {
+#ifdef IREE_UK_ENABLE_ASSERTS
+ const iree_uk_uint32_t allflags =
+ IREE_UK_FLAG_CONV_NCHWC_TYPE_MASK | IREE_UK_FLAG_CONV_NCHWC_ACCUMULATE |
+ IREE_UK_FLAG_CONV_NCHWC_ALLOW_GENERIC_FALLBACK_TILE_FUNCTION;
+ IREE_UK_ASSERT(!(params->flags & ~allflags));
+ const iree_uk_uint32_t type_bits =
+ params->flags & IREE_UK_FLAG_CONV_NCHWC_TYPE_MASK;
+ IREE_UK_ASSERT(type_bits > 0 && type_bits < IREE_UK_FLAG_CONV_NCHWC_TYPE_END);
+ // Bound shape dims to a narrower int32 range here. This can be
+ // relaxed later if needed.
+ IREE_UK_ASSERT(IREE_UK_VALUE_IN_UNSIGNED_INT_RANGE(params->N, 31));
+ IREE_UK_ASSERT(IREE_UK_VALUE_IN_UNSIGNED_INT_RANGE(params->OC_outer, 31));
+ IREE_UK_ASSERT(IREE_UK_VALUE_IN_UNSIGNED_INT_RANGE(params->OH, 31));
+ IREE_UK_ASSERT(IREE_UK_VALUE_IN_UNSIGNED_INT_RANGE(params->OW, 31));
+ IREE_UK_ASSERT(IREE_UK_VALUE_IN_UNSIGNED_INT_RANGE(params->IC_outer, 31));
+ IREE_UK_ASSERT(IREE_UK_VALUE_IN_UNSIGNED_INT_RANGE(params->FH, 31));
+ IREE_UK_ASSERT(IREE_UK_VALUE_IN_UNSIGNED_INT_RANGE(params->FW, 31));
+ // int32 is overkill for tile sizes; enforce int16 range for now.
+ IREE_UK_ASSERT(IREE_UK_VALUE_IN_UNSIGNED_INT_RANGE(params->k0, 15));
+ IREE_UK_ASSERT(IREE_UK_VALUE_IN_UNSIGNED_INT_RANGE(params->c0, 15));
+ IREE_UK_ASSERT(params->k0 > 0 && params->c0 > 0);
+ IREE_UK_ASSERT(params->stride_h > 0 && params->stride_w > 0);
+#endif
+}
+
+static bool iree_uk_conv_nchwc_early(
+ const iree_uk_conv_nchwc_params_t* params) {
+ if (params->N == 0 || params->OC_outer == 0 || params->OH == 0 ||
+ params->OW == 0) {
+ return true;
+ }
+ if ((params->IC_outer == 0 || params->FH == 0 || params->FW == 0) &&
+ (params->flags & IREE_UK_FLAG_CONV_NCHWC_ACCUMULATE)) {
+ return true;
+ }
+ return false;
+}
+
+// Shared outer driver for conv2d tile implementation.
+//
+// The tile function is called once per output panel of shape [OW_TILE, k0].
+// It receives:
+// - output_panel: pointer to output[n, oc_outer, oh, ow, 0] (k0 inner)
+// - input_panel: pointer to input[n, 0, oh x sH, ow x sW, 0]
+// - filter_panel: pointer to filter[oc_outer, 0, 0, 0, 0, 0]
+// - params: strides + shape info
+//
+// The tile function internally handles the reduction over (ic_outer, fh, fw,
+// c_inner) and honors the ACCUMULATE flag for the whole call.
+static void iree_uk_conv_nchwc_using_tile_func(
+ const iree_uk_conv_nchwc_params_t* params,
+ iree_uk_conv_nchwc_tile_func_t tile_func, iree_uk_int32_t OW_TILE) {
+ IREE_UK_ASSERT(OW_TILE > 0);
+ const iree_uk_index_t N = params->N;
+ const iree_uk_index_t OC_outer = params->OC_outer;
+ const iree_uk_index_t OH = params->OH;
+ const iree_uk_index_t OW = params->OW;
+ const iree_uk_int32_t k0 = params->k0;
+ const iree_uk_int32_t c0 = params->c0;
+ const iree_uk_int32_t stride_h = params->stride_h;
+ const iree_uk_int32_t stride_w = params->stride_w;
+ const iree_uk_index_t input_stride_n = params->input_stride_n;
+ const iree_uk_index_t input_stride_h = params->input_stride_h;
+ const iree_uk_index_t filter_stride_oc_outer = params->filter_stride_oc_outer;
+ const iree_uk_index_t output_stride_n = params->output_stride_n;
+ const iree_uk_index_t output_stride_oc_outer = params->output_stride_oc_outer;
+ const iree_uk_index_t output_stride_oh = params->output_stride_oh;
+
+ const iree_uk_conv_nchwc_type_t type = iree_uk_conv_nchwc_type(params->flags);
+ const iree_uk_index_t input_elem_size =
+ iree_uk_type_size(iree_uk_conv_nchwc_input_type(type));
+ const iree_uk_index_t filter_elem_size =
+ iree_uk_type_size(iree_uk_conv_nchwc_filter_type(type));
+ const iree_uk_index_t out_elem_size =
+ iree_uk_type_size(iree_uk_conv_nchwc_output_type(type));
+
+ const char* input_base = (const char*)params->input_buffer +
+ params->input_offset * input_elem_size;
+ const char* filter_base = (const char*)params->filter_buffer +
+ params->filter_offset * filter_elem_size;
+ char* output_base =
+ (char*)params->output_buffer + params->output_offset * out_elem_size;
+
+ iree_uk_conv_nchwc_params_t call_params = *params;
+
+ for (iree_uk_index_t n = 0; n < N; ++n) {
+ const char* in_n = input_base + n * input_stride_n * input_elem_size;
+ char* out_n = output_base + n * output_stride_n * out_elem_size;
+
+ for (iree_uk_index_t oc_o = 0; oc_o < OC_outer; ++oc_o) {
+ const char* f_panel =
+ filter_base + oc_o * filter_stride_oc_outer * filter_elem_size;
+ char* out_oc = out_n + oc_o * output_stride_oc_outer * out_elem_size;
+
+ for (iree_uk_index_t oh = 0; oh < OH; ++oh) {
+ char* out_oh = out_oc + oh * output_stride_oh * out_elem_size;
+ iree_uk_index_t ih = oh * stride_h;
+ const char* in_oh = in_n + ih * input_stride_h * input_elem_size;
+
+ for (iree_uk_index_t ow = 0; ow < OW; ow += OW_TILE) {
+ iree_uk_index_t ow_count = OW_TILE;
+ if (ow + OW_TILE > OW) {
+ ow_count = OW - ow;
+ }
+ char* out_panel = out_oh + ow * k0 * out_elem_size;
+ iree_uk_index_t iw = ow * stride_w;
+ const char* in_panel = in_oh + iw * c0 * input_elem_size;
+
+ // Active output width for this call; equals OW_TILE except on the
+ // right-edge tail.
+ call_params.ow_count = ow_count;
+ tile_func(out_panel, in_panel, f_panel, &call_params);
+ }
+ }
+ }
+ }
+}
+
+void iree_uk_conv_nchwc_p(const iree_uk_conv_nchwc_params_t* params) {
+ iree_uk_conv_nchwc_validate(params);
+
+ // Return early without selecting a tile function in trivial cases.
+ if (iree_uk_conv_nchwc_early(params)) return;
+
+ // Select a target specific tile function.
+ iree_uk_conv_nchwc_tile_selection_t sel =
+ iree_uk_conv_nchwc_select_tile_func_arch(params);
+
+ // If no target specific tile function is available, fall back to a generic
+ // one if allowed by the flags.
+ if (!sel.tile_func) {
+ if (params->flags &
+ IREE_UK_FLAG_CONV_NCHWC_ALLOW_GENERIC_FALLBACK_TILE_FUNCTION) {
+ sel = iree_uk_conv_nchwc_select_tile_func_generic(params);
+ } else {
+ IREE_UK_ASSERT(
+ 0 && "no target-specific tile function, and fallback not enabled.");
+ }
+ }
+ iree_uk_conv_nchwc_using_tile_func(params, sel.tile_func, sel.ow_tile);
+}
+
+IREE_UK_EXPORT void iree_uk_conv_nchwc(
+ const void* input_buffer, iree_uk_index_t input_offset,
+ iree_uk_index_t input_stride_n, iree_uk_index_t input_stride_ic_outer,
+ iree_uk_index_t input_stride_h, const void* filter_buffer,
+ iree_uk_index_t filter_offset, iree_uk_index_t filter_stride_oc_outer,
+ iree_uk_index_t filter_stride_ic_outer, iree_uk_index_t filter_stride_fh,
+ void* output_buffer, iree_uk_index_t output_offset,
+ iree_uk_index_t output_stride_n, iree_uk_index_t output_stride_oc_outer,
+ iree_uk_index_t output_stride_oh, iree_uk_index_t N,
+ iree_uk_index_t OC_outer, iree_uk_index_t OH, iree_uk_index_t OW,
+ iree_uk_index_t IC_outer, iree_uk_index_t FH, iree_uk_index_t FW,
+ iree_uk_int32_t k0, iree_uk_int32_t c0, iree_uk_int32_t stride_h,
+ iree_uk_int32_t stride_w, iree_uk_uint32_t flags,
+ const iree_uk_uint64_t* cpu_data) {
+ iree_uk_conv_nchwc_params_t params = {
+ .input_buffer = input_buffer,
+ .input_offset = input_offset,
+ .input_stride_n = input_stride_n,
+ .input_stride_ic_outer = input_stride_ic_outer,
+ .input_stride_h = input_stride_h,
+ .filter_buffer = filter_buffer,
+ .filter_offset = filter_offset,
+ .filter_stride_oc_outer = filter_stride_oc_outer,
+ .filter_stride_ic_outer = filter_stride_ic_outer,
+ .filter_stride_fh = filter_stride_fh,
+ .output_buffer = output_buffer,
+ .output_offset = output_offset,
+ .output_stride_n = output_stride_n,
+ .output_stride_oc_outer = output_stride_oc_outer,
+ .output_stride_oh = output_stride_oh,
+ .N = N,
+ .OC_outer = OC_outer,
+ .OH = OH,
+ .OW = OW,
+ .IC_outer = IC_outer,
+ .FH = FH,
+ .FW = FW,
+ .k0 = k0,
+ .c0 = c0,
+ .stride_h = stride_h,
+ .stride_w = stride_w,
+ .flags = flags,
+ .cpu_data = cpu_data,
+ };
+ iree_uk_conv_nchwc_p(¶ms);
+}
+
+iree_uk_uint32_t iree_uk_conv_nchwc_info_p(
+ const iree_uk_conv_nchwc_params_t* params) {
+ iree_uk_uint32_t result = 0;
+ if (iree_uk_conv_nchwc_select_tile_func_arch(params).tile_func) {
+ result |=
+ IREE_UK_FLAG_CONV_NCHWC_INFO_HAVE_ARCHITECTURE_SPECIFIC_TILE_FUNCTION;
+ }
+ return result;
+}
+
+IREE_UK_EXPORT iree_uk_uint32_t iree_uk_conv_nchwc_info(
+ iree_uk_int32_t k0, iree_uk_int32_t c0, iree_uk_uint32_t flags,
+ const iree_uk_uint64_t* cpu_data) {
+ iree_uk_conv_nchwc_params_t params = {
+ .k0 = k0, .c0 = c0, .flags = flags, .cpu_data = cpu_data};
+ return iree_uk_conv_nchwc_info_p(¶ms);
+}
diff --git a/runtime/src/iree/builtins/ukernel/conv_nchwc.h b/runtime/src/iree/builtins/ukernel/conv_nchwc.h
new file mode 100644
index 0000000..b4d8ff2
--- /dev/null
+++ b/runtime/src/iree/builtins/ukernel/conv_nchwc.h
@@ -0,0 +1,34 @@
+// Copyright 2026 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#ifndef IREE_BUILTINS_UKERNEL_CONV_NCHWC_H_
+#define IREE_BUILTINS_UKERNEL_CONV_NCHWC_H_
+
+#include "iree/builtins/ukernel/common.h"
+
+// `conv_nchwc` microkernel. Computes a data-tiled NCHWc 2D convolution.
+IREE_UK_EXPORT void iree_uk_conv_nchwc(
+ const void* input_buffer, iree_uk_index_t input_offset,
+ iree_uk_index_t input_stride_n, iree_uk_index_t input_stride_ic_outer,
+ iree_uk_index_t input_stride_h, const void* filter_buffer,
+ iree_uk_index_t filter_offset, iree_uk_index_t filter_stride_oc_outer,
+ iree_uk_index_t filter_stride_ic_outer, iree_uk_index_t filter_stride_fh,
+ void* output_buffer, iree_uk_index_t output_offset,
+ iree_uk_index_t output_stride_n, iree_uk_index_t output_stride_oc_outer,
+ iree_uk_index_t output_stride_oh, iree_uk_index_t N,
+ iree_uk_index_t OC_outer, iree_uk_index_t OH, iree_uk_index_t OW,
+ iree_uk_index_t IC_outer, iree_uk_index_t FH, iree_uk_index_t FW,
+ iree_uk_int32_t k0, iree_uk_int32_t c0, iree_uk_int32_t stride_h,
+ iree_uk_int32_t stride_w, iree_uk_uint32_t flags,
+ const iree_uk_uint64_t* cpu_data);
+
+// Returns a bit-field of information about how a conv_nchwc with the given
+// tile shape and flags would run on the current target.
+IREE_UK_EXPORT iree_uk_uint32_t iree_uk_conv_nchwc_info(
+ iree_uk_int32_t k0, iree_uk_int32_t c0, iree_uk_uint32_t flags,
+ const iree_uk_uint64_t* cpu_data);
+
+#endif // IREE_BUILTINS_UKERNEL_CONV_NCHWC_H_
diff --git a/runtime/src/iree/builtins/ukernel/conv_nchwc_internal.h b/runtime/src/iree/builtins/ukernel/conv_nchwc_internal.h
new file mode 100644
index 0000000..f12d3d4
--- /dev/null
+++ b/runtime/src/iree/builtins/ukernel/conv_nchwc_internal.h
@@ -0,0 +1,127 @@
+// Copyright 2026 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#ifndef IREE_BUILTINS_UKERNEL_CONV_NCHWC_INTERNAL_H_
+#define IREE_BUILTINS_UKERNEL_CONV_NCHWC_INTERNAL_H_
+
+#include "iree/builtins/ukernel/conv_nchwc.h"
+
+// Internal params struct passed between implementation functions.
+typedef struct iree_uk_conv_nchwc_params_t {
+ const void* input_buffer;
+ iree_uk_index_t input_offset;
+ iree_uk_index_t input_stride_n;
+ iree_uk_index_t input_stride_ic_outer;
+ iree_uk_index_t input_stride_h;
+ const void* filter_buffer;
+ iree_uk_index_t filter_offset;
+ iree_uk_index_t filter_stride_oc_outer;
+ iree_uk_index_t filter_stride_ic_outer;
+ iree_uk_index_t filter_stride_fh;
+ void* output_buffer;
+ iree_uk_index_t output_offset;
+ iree_uk_index_t output_stride_n;
+ iree_uk_index_t output_stride_oc_outer;
+ iree_uk_index_t output_stride_oh;
+ iree_uk_index_t N;
+ iree_uk_index_t OC_outer;
+ iree_uk_index_t OH;
+ iree_uk_index_t OW;
+ iree_uk_index_t IC_outer;
+ iree_uk_index_t FH;
+ iree_uk_index_t FW;
+ iree_uk_int32_t k0; // inner OC block
+ iree_uk_int32_t c0; // inner IC block
+ iree_uk_int32_t stride_h;
+ iree_uk_int32_t stride_w;
+ iree_uk_uint32_t flags;
+ // ow_count: how many output columns this particular tile call should produce.
+ // Equals the selected tile's ow_tile for full tiles; smaller only on the
+ // partial tail, as set by the driver.
+ iree_uk_index_t ow_count;
+ const iree_uk_uint64_t* cpu_data;
+} iree_uk_conv_nchwc_params_t;
+
+// Equivalent to the public iree_uk_conv_nchwc entry point but taking the params
+// struct.
+void iree_uk_conv_nchwc_p(const iree_uk_conv_nchwc_params_t* params);
+
+// Equivalent to the public iree_uk_conv_nchwc_info but operating on the params
+// struct. Only the struct fields corresponding to iree_uk_conv_nchwc_info
+// params are used.
+iree_uk_uint32_t iree_uk_conv_nchwc_info_p(
+ const iree_uk_conv_nchwc_params_t* params);
+
+// TODO: extend to match mmt4d's coverage (S8S8S32, S16/U4 variants, F16,
+// BF16) once those paths are exercised by the compiler-side encoding
+// resolver. For now this ukernel is F32-only.
+typedef enum iree_uk_conv_nchwc_type_t {
+ iree_uk_conv_nchwc_type_f32f32f32 =
+ IREE_UK_TIE_3_TYPES_LITERAL(FLOAT_32, FLOAT_32, FLOAT_32),
+} iree_uk_conv_nchwc_type_t;
+
+static inline iree_uk_conv_nchwc_type_t iree_uk_conv_nchwc_type(
+ iree_uk_uint32_t flags) {
+ switch (flags & IREE_UK_FLAG_CONV_NCHWC_TYPE_MASK) {
+ case IREE_UK_FLAG_CONV_NCHWC_TYPE_F32F32F32:
+ return iree_uk_conv_nchwc_type_f32f32f32;
+ default:
+ // Work around a LLVM/riscv32 miscompile. Without the unreachable here,
+ // returning (iree_uk_conv_nchwc_type_t)0 causes this whole switch
+ // statement to be miscompiled by LLVM/riscv32 as if it were UB, as the 0
+ // was passed to `iree_uk_type_bit_count(x)`, which evaluates to
+ // `1<<(x - 3)`, which is UB if x<3.
+#if defined(IREE_UK_COMPILER_CLANG) && defined(IREE_UK_ARCH_RISCV_32)
+ __builtin_unreachable();
+#endif
+ // Shouldn't happen, validated earlier.
+ return (iree_uk_conv_nchwc_type_t)0;
+ }
+}
+
+static inline iree_uk_type_t iree_uk_conv_nchwc_input_type(
+ iree_uk_conv_nchwc_type_t type) {
+ return iree_uk_untie_type(0, type);
+}
+
+static inline iree_uk_type_t iree_uk_conv_nchwc_filter_type(
+ iree_uk_conv_nchwc_type_t type) {
+ return iree_uk_untie_type(1, type);
+}
+
+static inline iree_uk_type_t iree_uk_conv_nchwc_output_type(
+ iree_uk_conv_nchwc_type_t type) {
+ return iree_uk_untie_type(2, type);
+}
+
+// Function pointer that computes one (ow_count x k0) output panel.
+// The tile function performs the full reduction over (ic_outer, fh, fw,
+// c_inner) internally. The driver calls it once per (n, oc_outer, oh, ow-tile)
+// output panel. When the reduction range is empty (IC_outer/FH/FW == 0), the
+// driver still calls the tile function in the non-accumulate case, and the tile
+// function must write zeros to the output panel.
+typedef void (*iree_uk_conv_nchwc_tile_func_t)(
+ void* IREE_UK_RESTRICT output_panel,
+ const void* IREE_UK_RESTRICT input_panel,
+ const void* IREE_UK_RESTRICT filter_panel,
+ const iree_uk_conv_nchwc_params_t* params);
+
+// A selected tile function plus the output-width tile size it was compiled
+// for.
+typedef struct iree_uk_conv_nchwc_tile_selection_t {
+ iree_uk_conv_nchwc_tile_func_t tile_func; // null if unavailable
+ iree_uk_int32_t ow_tile;
+} iree_uk_conv_nchwc_tile_selection_t;
+
+// Architecture-specific implementation, or generic fallback returning null
+iree_uk_conv_nchwc_tile_selection_t iree_uk_conv_nchwc_select_tile_func_arch(
+ const iree_uk_conv_nchwc_params_t* params);
+
+// Generic fallback.
+iree_uk_conv_nchwc_tile_selection_t iree_uk_conv_nchwc_select_tile_func_generic(
+ const iree_uk_conv_nchwc_params_t* params);
+
+#endif // IREE_BUILTINS_UKERNEL_CONV_NCHWC_INTERNAL_H_
diff --git a/runtime/src/iree/builtins/ukernel/conv_nchwc_tile_generic.c b/runtime/src/iree/builtins/ukernel/conv_nchwc_tile_generic.c
new file mode 100644
index 0000000..b6bf758
--- /dev/null
+++ b/runtime/src/iree/builtins/ukernel/conv_nchwc_tile_generic.c
@@ -0,0 +1,77 @@
+// Copyright 2026 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/builtins/ukernel/conv_nchwc_internal.h"
+#include "iree/builtins/ukernel/exported_bits.h"
+
+// Generic scalar implementation of the conv_nchwc tile function for f32.
+//
+// Computes one output panel of shape [ow_count, k0] by iterating over
+// (ic_outer, fh, fw, c_inner) and accumulating. Used as a reference kernel
+// for correctness checking and as a portable fallback when no ISA-specific
+// tile function is available.
+static void iree_uk_conv_nchwc_tile_f32f32f32_generic(
+ void* IREE_UK_RESTRICT output_panel,
+ const void* IREE_UK_RESTRICT input_panel,
+ const void* IREE_UK_RESTRICT filter_panel,
+ const iree_uk_conv_nchwc_params_t* params) {
+ float* out_ptr = (float*)output_panel;
+ const float* in_ptr = (const float*)input_panel;
+ const float* f_ptr = (const float*)filter_panel;
+
+ const iree_uk_index_t ow_count = params->ow_count;
+ const iree_uk_index_t IC_outer = params->IC_outer;
+ const iree_uk_index_t FH = params->FH;
+ const iree_uk_index_t FW = params->FW;
+ const iree_uk_int32_t k0 = params->k0;
+ const iree_uk_int32_t c0 = params->c0;
+ const iree_uk_int32_t stride_w = params->stride_w;
+
+ const iree_uk_index_t in_stride_w = c0;
+ const iree_uk_index_t in_stride_h = params->input_stride_h;
+ const iree_uk_index_t in_stride_ic_outer = params->input_stride_ic_outer;
+
+ const iree_uk_index_t f_stride_fw = (iree_uk_index_t)c0 * k0;
+ const iree_uk_index_t f_stride_fh = params->filter_stride_fh;
+ const iree_uk_index_t f_stride_ic_outer = params->filter_stride_ic_outer;
+
+ for (iree_uk_index_t i = 0; i < ow_count; ++i) {
+ for (iree_uk_int32_t ko = 0; ko < k0; ++ko) {
+ float acc = (params->flags & IREE_UK_FLAG_CONV_NCHWC_ACCUMULATE)
+ ? out_ptr[i * k0 + ko]
+ : 0.f;
+ for (iree_uk_index_t ic_o = 0; ic_o < IC_outer; ++ic_o) {
+ for (iree_uk_index_t fh = 0; fh < FH; ++fh) {
+ for (iree_uk_index_t fw = 0; fw < FW; ++fw) {
+ for (iree_uk_int32_t ci = 0; ci < c0; ++ci) {
+ float in_val =
+ in_ptr[ic_o * in_stride_ic_outer + fh * in_stride_h +
+ (i * stride_w + fw) * in_stride_w + ci];
+ float f_val = f_ptr[ic_o * f_stride_ic_outer + fh * f_stride_fh +
+ fw * f_stride_fw + ci * k0 + ko];
+ acc += in_val * f_val;
+ }
+ }
+ }
+ }
+ out_ptr[i * k0 + ko] = acc;
+ }
+ }
+}
+
+iree_uk_conv_nchwc_tile_selection_t iree_uk_conv_nchwc_select_tile_func_generic(
+ const iree_uk_conv_nchwc_params_t* params) {
+ iree_uk_conv_nchwc_tile_selection_t selection = {0};
+ switch (iree_uk_conv_nchwc_type(params->flags)) {
+ case iree_uk_conv_nchwc_type_f32f32f32:
+ selection.tile_func = iree_uk_conv_nchwc_tile_f32f32f32_generic;
+ selection.ow_tile = 16;
+ break;
+ default:
+ break;
+ }
+ return selection;
+}
diff --git a/runtime/src/iree/builtins/ukernel/exported_bits.h b/runtime/src/iree/builtins/ukernel/exported_bits.h
index aa25b0d..571ceef 100644
--- a/runtime/src/iree/builtins/ukernel/exported_bits.h
+++ b/runtime/src/iree/builtins/ukernel/exported_bits.h
@@ -105,4 +105,22 @@
#define IREE_UK_FLAG_QUERY_TILE_SIZES_OPERATION_MATMUL_BF16BF16F32 0x0500
#define IREE_UK_FLAG_QUERY_TILE_SIZES_OPERATION_MATMUL_BF16BF16BF16 0x0600
+//===----------------------------------------------------------------------===//
+// conv_nchwc
+//===----------------------------------------------------------------------===//
+
+// type enum (lhs = input, rhs = filter, out = output)
+#define IREE_UK_FLAG_CONV_NCHWC_TYPE_MASK 0xFF
+#define IREE_UK_FLAG_CONV_NCHWC_TYPE_NONE 0x00
+#define IREE_UK_FLAG_CONV_NCHWC_TYPE_F32F32F32 0x01
+#define IREE_UK_FLAG_CONV_NCHWC_TYPE_END 0x02
+
+// bit flags
+#define IREE_UK_FLAG_CONV_NCHWC_ACCUMULATE 0x100
+#define IREE_UK_FLAG_CONV_NCHWC_ALLOW_GENERIC_FALLBACK_TILE_FUNCTION 0x200
+
+// info bit flags returned by iree_uk_conv_nchwc_info.
+#define IREE_UK_FLAG_CONV_NCHWC_INFO_HAVE_ARCHITECTURE_SPECIFIC_TILE_FUNCTION \
+ 0x1
+
#endif // IREE_BUILTINS_UKERNEL_EXPORTED_BITS_H_
diff --git a/runtime/src/iree/builtins/ukernel/fallback.c b/runtime/src/iree/builtins/ukernel/fallback.c
index eaee5b2..779e7da 100644
--- a/runtime/src/iree/builtins/ukernel/fallback.c
+++ b/runtime/src/iree/builtins/ukernel/fallback.c
@@ -4,6 +4,7 @@
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+#include "iree/builtins/ukernel/conv_nchwc_internal.h"
#include "iree/builtins/ukernel/mmt4d_internal.h"
#include "iree/builtins/ukernel/pack_internal.h"
#include "iree/builtins/ukernel/query_tile_sizes_internal.h"
@@ -29,3 +30,9 @@
iree_uk_matmul_tile_sizes_t* out_matmul_tile_sizes) {
return false;
}
+
+iree_uk_conv_nchwc_tile_selection_t iree_uk_conv_nchwc_select_tile_func_arch(
+ const iree_uk_conv_nchwc_params_t* params) {
+ iree_uk_conv_nchwc_tile_selection_t selection = {0};
+ return selection;
+}
diff --git a/runtime/src/iree/builtins/ukernel/tools/BUILD.bazel b/runtime/src/iree/builtins/ukernel/tools/BUILD.bazel
index 2f5d41c..d543288 100644
--- a/runtime/src/iree/builtins/ukernel/tools/BUILD.bazel
+++ b/runtime/src/iree/builtins/ukernel/tools/BUILD.bazel
@@ -104,6 +104,33 @@
],
)
+iree_runtime_cc_test(
+ name = "conv_nchwc_test",
+ srcs = ["conv_nchwc_test.c"],
+ deps = [
+ ":test",
+ ":util",
+ "//runtime/src/iree/base",
+ "//runtime/src/iree/base/internal",
+ "//runtime/src/iree/builtins/ukernel",
+ "//runtime/src/iree/builtins/ukernel:internal_headers",
+ ],
+)
+
+cc_binary_benchmark(
+ name = "conv_nchwc_benchmark",
+ srcs = ["conv_nchwc_benchmark.c"],
+ deps = [
+ ":benchmark",
+ ":util",
+ "//runtime/src/iree/base",
+ "//runtime/src/iree/base/tooling:flags",
+ "//runtime/src/iree/builtins/ukernel",
+ "//runtime/src/iree/builtins/ukernel:internal_headers",
+ "//runtime/src/iree/testing:benchmark",
+ ],
+)
+
cc_binary_benchmark(
name = "pack_benchmark",
srcs = ["pack_benchmark.c"],
diff --git a/runtime/src/iree/builtins/ukernel/tools/CMakeLists.txt b/runtime/src/iree/builtins/ukernel/tools/CMakeLists.txt
index 490d0b7..e9700a5 100644
--- a/runtime/src/iree/builtins/ukernel/tools/CMakeLists.txt
+++ b/runtime/src/iree/builtins/ukernel/tools/CMakeLists.txt
@@ -117,6 +117,36 @@
iree::builtins::ukernel::internal_headers
)
+iree_cc_test(
+ NAME
+ conv_nchwc_test
+ SRCS
+ "conv_nchwc_test.c"
+ DEPS
+ ::test
+ ::util
+ iree::base
+ iree::base::internal
+ iree::builtins::ukernel
+ iree::builtins::ukernel::internal_headers
+)
+
+iree_cc_binary_benchmark(
+ NAME
+ conv_nchwc_benchmark
+ SRCS
+ "conv_nchwc_benchmark.c"
+ DEPS
+ ::benchmark
+ ::util
+ iree::base
+ iree::base::tooling::flags
+ iree::builtins::ukernel
+ iree::builtins::ukernel::internal_headers
+ iree::testing::benchmark
+ TESTONLY
+)
+
iree_cc_binary_benchmark(
NAME
pack_benchmark
diff --git a/runtime/src/iree/builtins/ukernel/tools/conv_nchwc_benchmark.c b/runtime/src/iree/builtins/ukernel/tools/conv_nchwc_benchmark.c
new file mode 100644
index 0000000..a36ab9e
--- /dev/null
+++ b/runtime/src/iree/builtins/ukernel/tools/conv_nchwc_benchmark.c
@@ -0,0 +1,137 @@
+// Copyright 2026 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include <stdio.h>
+
+#include "iree/base/api.h"
+#include "iree/base/tooling/flags.h"
+#include "iree/builtins/ukernel/api.h"
+#include "iree/builtins/ukernel/conv_nchwc.h"
+#include "iree/builtins/ukernel/conv_nchwc_internal.h"
+#include "iree/builtins/ukernel/exported_bits.h"
+#include "iree/builtins/ukernel/tools/benchmark.h"
+#include "iree/builtins/ukernel/tools/util.h"
+
+IREE_FLAG(int32_t, n_size, 1, "Batch (N) of the conv.");
+IREE_FLAG(int32_t, oc_size, 8, "Output-channel tiles (OC / k0).");
+IREE_FLAG(int32_t, ic_size, 8, "Input-channel tiles (IC / c0).");
+IREE_FLAG(int32_t, oh_size, 28, "Output height (OH).");
+IREE_FLAG(int32_t, ow_size, 28, "Output width (OW).");
+IREE_FLAG(int32_t, fh_size, 3, "Filter height (FH).");
+IREE_FLAG(int32_t, fw_size, 3, "Filter width (FW).");
+IREE_FLAG(bool, accumulate, false,
+ "Whether the kernel should accumulate into the existing output, or "
+ "overwrite it.");
+
+static iree_status_t iree_uk_benchmark_conv_nchwc(
+ const iree_benchmark_def_t* benchmark_def,
+ iree_benchmark_state_t* benchmark_state) {
+ const iree_uk_benchmark_user_data_t* user_data = benchmark_def->user_data;
+ const iree_uk_conv_nchwc_params_t* src_params =
+ iree_uk_benchmark_params(user_data);
+ iree_uk_conv_nchwc_params_t params;
+ memcpy(¶ms, src_params, sizeof params);
+ params.cpu_data = iree_uk_benchmark_cpu_data(user_data);
+ if (FLAG_accumulate) params.flags |= IREE_UK_FLAG_CONV_NCHWC_ACCUMULATE;
+ params.N = FLAG_n_size;
+ params.OC_outer = FLAG_oc_size;
+ params.IC_outer = FLAG_ic_size;
+ params.OH = FLAG_oh_size;
+ params.OW = FLAG_ow_size;
+ params.FH = FLAG_fh_size;
+ params.FW = FLAG_fw_size;
+
+ iree_uk_index_t H = (params.OH - 1) * params.stride_h + params.FH;
+ iree_uk_index_t W = (params.OW - 1) * params.stride_w + params.FW;
+ params.input_stride_h = W * params.c0;
+ params.input_stride_ic_outer = H * params.input_stride_h;
+ params.input_stride_n = params.IC_outer * params.input_stride_ic_outer;
+ iree_uk_index_t filter_stride_fw = (iree_uk_index_t)params.c0 * params.k0;
+ params.filter_stride_fh = params.FW * filter_stride_fw;
+ params.filter_stride_ic_outer = params.FH * params.filter_stride_fh;
+ params.filter_stride_oc_outer =
+ params.IC_outer * params.filter_stride_ic_outer;
+ params.output_stride_oh = params.OW * params.k0;
+ params.output_stride_oc_outer = params.OH * params.output_stride_oh;
+ params.output_stride_n = params.OC_outer * params.output_stride_oc_outer;
+
+ iree_uk_conv_nchwc_type_t type = iree_uk_conv_nchwc_type(params.flags);
+ iree_uk_type_t in_type = iree_uk_conv_nchwc_input_type(type);
+ iree_uk_type_t f_type = iree_uk_conv_nchwc_filter_type(type);
+ iree_uk_type_t out_type = iree_uk_conv_nchwc_output_type(type);
+ iree_uk_index_t in_size =
+ iree_uk_2d_buffer_length(in_type, params.N, params.input_stride_n);
+ iree_uk_index_t f_size = iree_uk_2d_buffer_length(
+ f_type, params.OC_outer, params.filter_stride_oc_outer);
+ iree_uk_index_t out_size =
+ iree_uk_2d_buffer_length(out_type, params.N, params.output_stride_n);
+ void* in_buffer = malloc(in_size);
+ void* f_buffer = malloc(f_size);
+ void* out_buffer = malloc(out_size);
+ iree_uk_random_engine_t* engine = iree_uk_benchmark_random_engine(user_data);
+ iree_uk_write_random_buffer(in_buffer, in_size, in_type, engine);
+ iree_uk_write_random_buffer(f_buffer, f_size, f_type, engine);
+ iree_uk_write_random_buffer(out_buffer, out_size, out_type, engine);
+ params.input_buffer = in_buffer;
+ params.filter_buffer = f_buffer;
+ params.output_buffer = out_buffer;
+
+ int64_t total_iterations = 0;
+ int64_t batch_count = 1;
+ while (iree_benchmark_keep_running(benchmark_state, batch_count)) {
+ for (int i = 0; i < batch_count; ++i) {
+ iree_uk_conv_nchwc_p(¶ms);
+ }
+ total_iterations += batch_count;
+ batch_count *= 2;
+ }
+ // FLOPs: 2 (multiply-add) per (output element x reduction element).
+ iree_benchmark_set_items_processed(
+ benchmark_state, total_iterations * 2 * params.N * params.OC_outer *
+ params.k0 * params.OH * params.OW * params.IC_outer *
+ params.c0 * params.FH * params.FW);
+ free(in_buffer);
+ free(f_buffer);
+ free(out_buffer);
+ return iree_ok_status();
+}
+
+static void iree_uk_benchmark_register_conv_nchwc(iree_uk_uint32_t flags,
+ int k0, int c0,
+ const char* cpu_features) {
+ char type_str[32];
+ iree_uk_conv_nchwc_type_t type = iree_uk_conv_nchwc_type(flags);
+ iree_uk_type_triple_str(type_str, sizeof type_str, type);
+ char name[128];
+ iree_snprintf(name, sizeof name, "conv_nchwc_%s_tile_%dx%d", type_str, k0,
+ c0);
+ iree_uk_conv_nchwc_params_t params = {
+ .flags =
+ flags | IREE_UK_FLAG_CONV_NCHWC_ALLOW_GENERIC_FALLBACK_TILE_FUNCTION,
+ .k0 = k0,
+ .c0 = c0,
+ .stride_h = 1,
+ .stride_w = 1};
+ iree_uk_benchmark_register(name, iree_uk_benchmark_conv_nchwc, ¶ms,
+ sizeof params, cpu_features);
+}
+
+int main(int argc, char** argv) {
+ iree_flags_set_usage("conv_nchwc_benchmark", "");
+
+ iree_flags_parse_checked(IREE_FLAGS_PARSE_MODE_UNDEFINED_OK, &argc, &argv);
+ iree_uk_benchmark_initialize(&argc, argv);
+
+#if defined(IREE_ARCH_X86_64)
+ iree_uk_benchmark_register_conv_nchwc(IREE_UK_FLAG_CONV_NCHWC_TYPE_F32F32F32,
+ 16, 16, "avx512_base");
+#endif
+ // Generic fallback (no CPU features), for comparison against the arch tile.
+ iree_uk_benchmark_register_conv_nchwc(IREE_UK_FLAG_CONV_NCHWC_TYPE_F32F32F32,
+ 16, 16, "");
+
+ iree_uk_benchmark_run_and_cleanup();
+}
diff --git a/runtime/src/iree/builtins/ukernel/tools/conv_nchwc_test.c b/runtime/src/iree/builtins/ukernel/tools/conv_nchwc_test.c
new file mode 100644
index 0000000..c812b6a
--- /dev/null
+++ b/runtime/src/iree/builtins/ukernel/tools/conv_nchwc_test.c
@@ -0,0 +1,312 @@
+// Copyright 2026 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/base/api.h"
+#include "iree/builtins/ukernel/api.h"
+#include "iree/builtins/ukernel/conv_nchwc_internal.h"
+#include "iree/builtins/ukernel/exported_bits.h"
+#include "iree/builtins/ukernel/tools/test.h"
+#include "iree/builtins/ukernel/tools/util.h"
+
+// Inner reduction (ic_o, fh, fw, ci) for one output element. The outer driver
+// below computes `in_window` (pointer to input[n, 0, oh*sH, ow*sW, 0]) and
+// `f_slice` (pointer to filter[oc_o, 0, 0, 0, 0, ko]) and dispatches into here.
+static void iree_uk_conv_nchwc_reference_innerloop_f32f32f32(
+ float* out_ptr, const float* in_window, const float* f_slice,
+ const iree_uk_conv_nchwc_params_t* params) {
+ float acc =
+ params->flags & IREE_UK_FLAG_CONV_NCHWC_ACCUMULATE ? *out_ptr : 0.f;
+ const iree_uk_index_t c0 = params->c0;
+ const iree_uk_index_t k0 = params->k0;
+ for (iree_uk_index_t ic_o = 0; ic_o < params->IC_outer; ++ic_o) {
+ const float* in_ic = in_window + ic_o * params->input_stride_ic_outer;
+ const float* f_ic = f_slice + ic_o * params->filter_stride_ic_outer;
+ for (iree_uk_index_t fh = 0; fh < params->FH; ++fh) {
+ const float* in_fh = in_ic + fh * params->input_stride_h;
+ const float* f_fh = f_ic + fh * params->filter_stride_fh;
+ for (iree_uk_index_t fw = 0; fw < params->FW; ++fw) {
+ const float* f_fw = f_fh + fw * (c0 * k0);
+ for (iree_uk_index_t ci = 0; ci < c0; ++ci) {
+ acc += in_fh[fw * c0 + ci] * f_fw[ci * k0];
+ }
+ }
+ }
+ }
+ *out_ptr = acc;
+}
+
+static void iree_uk_conv_nchwc_reference(
+ const iree_uk_conv_nchwc_params_t* params) {
+ iree_uk_conv_nchwc_type_t type = iree_uk_conv_nchwc_type(params->flags);
+ iree_uk_index_t in_elem_size =
+ iree_uk_type_size(iree_uk_conv_nchwc_input_type(type));
+ iree_uk_index_t f_elem_size =
+ iree_uk_type_size(iree_uk_conv_nchwc_filter_type(type));
+ iree_uk_index_t out_elem_size =
+ iree_uk_type_size(iree_uk_conv_nchwc_output_type(type));
+ const char* in_base =
+ (const char*)params->input_buffer + params->input_offset * in_elem_size;
+ const char* f_base =
+ (const char*)params->filter_buffer + params->filter_offset * f_elem_size;
+ char* out_base =
+ (char*)params->output_buffer + params->output_offset * out_elem_size;
+
+ for (iree_uk_index_t n = 0; n < params->N; ++n) {
+ for (iree_uk_index_t oc_o = 0; oc_o < params->OC_outer; ++oc_o) {
+ for (iree_uk_index_t oh = 0; oh < params->OH; ++oh) {
+ for (iree_uk_index_t ow = 0; ow < params->OW; ++ow) {
+ const char* in_window =
+ in_base + (n * params->input_stride_n +
+ oh * params->stride_h * params->input_stride_h +
+ ow * params->stride_w * params->c0) *
+ in_elem_size;
+ char* out_pixel =
+ out_base + (n * params->output_stride_n +
+ oc_o * params->output_stride_oc_outer +
+ oh * params->output_stride_oh + ow * params->k0) *
+ out_elem_size;
+ for (iree_uk_index_t ko = 0; ko < params->k0; ++ko) {
+ const char* f_slice =
+ f_base +
+ (oc_o * params->filter_stride_oc_outer + ko) * f_elem_size;
+ void* out_ptr = out_pixel + ko * out_elem_size;
+ switch (params->flags & IREE_UK_FLAG_CONV_NCHWC_TYPE_MASK) {
+ case IREE_UK_FLAG_CONV_NCHWC_TYPE_F32F32F32:
+ iree_uk_conv_nchwc_reference_innerloop_f32f32f32(
+ (float*)out_ptr, (const float*)in_window,
+ (const float*)f_slice, params);
+ break;
+ default:
+ IREE_UK_ASSERT(false && "unhandled type");
+ }
+ }
+ }
+ }
+ }
+ }
+}
+
+// Strides/Offsets are element counts. For a sub-byte type (e.g., i4), an
+// odd count is a fractional number of bytes which isn't byte addressable.
+// Round up to the next whole-byte count.
+static iree_uk_index_t iree_uk_test_round_up_to_ensure_multiple_of_8_bits(
+ iree_uk_index_t index, iree_uk_type_t type) {
+ while ((index << iree_uk_type_bit_count_log2(type)) & 7) ++index;
+ return index;
+}
+
+// Randomly make the strides either tight or not to exercise all cases.
+static iree_uk_index_t iree_uk_test_random_stride(
+ iree_uk_index_t min_stride, iree_uk_type_t type,
+ iree_uk_random_engine_t* engine) {
+ iree_uk_index_t stride = min_stride + iree_uk_random_engine_get_0_1(engine);
+ return iree_uk_test_round_up_to_ensure_multiple_of_8_bits(stride, type);
+}
+
+// Randomly make the offsets either tight or not to exercise all cases.
+static iree_uk_index_t iree_uk_test_random_offset(
+ iree_uk_type_t type, iree_uk_random_engine_t* engine) {
+ iree_uk_index_t offset = iree_uk_random_engine_get_0_1(engine);
+ return iree_uk_test_round_up_to_ensure_multiple_of_8_bits(offset, type);
+}
+
+static void iree_uk_test_conv_nchwc_for_shape_params(
+ iree_uk_test_t* test, const iree_uk_conv_nchwc_params_t* src_params) {
+ iree_uk_conv_nchwc_params_t params;
+ memcpy(¶ms, src_params, sizeof params);
+ iree_uk_conv_nchwc_type_t type = iree_uk_conv_nchwc_type(params.flags);
+ iree_uk_type_t in_type = iree_uk_conv_nchwc_input_type(type);
+ iree_uk_type_t f_type = iree_uk_conv_nchwc_filter_type(type);
+ iree_uk_type_t out_type = iree_uk_conv_nchwc_output_type(type);
+ iree_uk_random_engine_t* engine = iree_uk_test_random_engine(test);
+
+ // Input H/W are derived from output shape, filter shape, and stride.
+ iree_uk_index_t H = (params.OH - 1) * params.stride_h + params.FH;
+ iree_uk_index_t W = (params.OW - 1) * params.stride_w + params.FW;
+
+ // Populate the outer strides; innermost dims (W, OW, FW, c0, k0) stay
+ // pack-contiguous. We need these before sizing buffers below.
+ params.input_stride_h =
+ iree_uk_test_random_stride(W * params.c0, in_type, engine);
+ params.input_stride_ic_outer =
+ iree_uk_test_random_stride(H * params.input_stride_h, in_type, engine);
+ params.input_stride_n = iree_uk_test_random_stride(
+ params.IC_outer * params.input_stride_ic_outer, in_type, engine);
+
+ iree_uk_index_t filter_stride_fw = (iree_uk_index_t)params.c0 * params.k0;
+ params.filter_stride_fh =
+ iree_uk_test_random_stride(params.FW * filter_stride_fw, f_type, engine);
+ params.filter_stride_ic_outer = iree_uk_test_random_stride(
+ params.FH * params.filter_stride_fh, f_type, engine);
+ params.filter_stride_oc_outer = iree_uk_test_random_stride(
+ params.IC_outer * params.filter_stride_ic_outer, f_type, engine);
+
+ params.output_stride_oh =
+ iree_uk_test_random_stride(params.OW * params.k0, out_type, engine);
+ params.output_stride_oc_outer = iree_uk_test_random_stride(
+ params.OH * params.output_stride_oh, out_type, engine);
+ params.output_stride_n = iree_uk_test_random_stride(
+ params.OC_outer * params.output_stride_oc_outer, out_type, engine);
+
+ iree_uk_index_t in_buffer_size =
+ iree_uk_2d_buffer_length(in_type, params.N, params.input_stride_n);
+ iree_uk_index_t f_buffer_size = iree_uk_2d_buffer_length(
+ f_type, params.OC_outer, params.filter_stride_oc_outer);
+ iree_uk_index_t out_buffer_size =
+ iree_uk_2d_buffer_length(out_type, params.N, params.output_stride_n);
+ void* in_buffer = malloc(in_buffer_size);
+ void* f_buffer = malloc(f_buffer_size);
+ iree_uk_write_random_buffer(in_buffer, in_buffer_size, in_type, engine);
+ iree_uk_write_random_buffer(f_buffer, f_buffer_size, f_type, engine);
+
+ // Random offsets + pointer-shift trick: shift `buffer` back by `offset`, so
+ // the kernel's `buffer + offset` lands at the actual allocation. Exercises
+ // the kernel's offset arithmetic without inflating the allocation.
+ params.input_offset = iree_uk_test_random_offset(in_type, engine);
+ params.filter_offset = iree_uk_test_random_offset(f_type, engine);
+ params.output_offset = iree_uk_test_random_offset(out_type, engine);
+ params.input_buffer =
+ (const char*)in_buffer -
+ iree_uk_bits_to_bytes_exact(params.input_offset
+ << iree_uk_type_bit_count_log2(in_type));
+ params.filter_buffer =
+ (const char*)f_buffer -
+ iree_uk_bits_to_bytes_exact(params.filter_offset
+ << iree_uk_type_bit_count_log2(f_type));
+
+ void* init_out_buffer = malloc(out_buffer_size);
+ iree_uk_write_random_buffer(init_out_buffer, out_buffer_size, out_type,
+ engine);
+ void* reference_out_buffer = malloc(out_buffer_size);
+ void* actual_out_buffer = malloc(out_buffer_size);
+ memcpy(reference_out_buffer, init_out_buffer, out_buffer_size);
+ memcpy(actual_out_buffer, init_out_buffer, out_buffer_size);
+
+ iree_uk_conv_nchwc_params_t reference_params;
+ memcpy(&reference_params, ¶ms, sizeof params);
+ reference_params.output_buffer =
+ (char*)reference_out_buffer -
+ iree_uk_bits_to_bytes_exact(params.output_offset
+ << iree_uk_type_bit_count_log2(out_type));
+
+ iree_uk_conv_nchwc_params_t actual_params;
+ memcpy(&actual_params, ¶ms, sizeof params);
+ actual_params.output_buffer =
+ (char*)actual_out_buffer -
+ iree_uk_bits_to_bytes_exact(params.output_offset
+ << iree_uk_type_bit_count_log2(out_type));
+
+ iree_uk_conv_nchwc_reference(&reference_params);
+ iree_uk_conv_nchwc_p(&actual_params);
+
+ // Exact compare: small-integer random fills keep every intermediate FMA
+ // exactly representable, so reduction order is immaterial.
+ if (memcmp(actual_out_buffer, reference_out_buffer, out_buffer_size)) {
+ IREE_UK_TEST_FAIL(test);
+ }
+
+ free(in_buffer);
+ free(f_buffer);
+ free(init_out_buffer);
+ free(reference_out_buffer);
+ free(actual_out_buffer);
+}
+
+static void iree_uk_test_conv_nchwc_for_tile_params(iree_uk_test_t* test,
+ const void* src_params) {
+ typedef struct conv_shape_t {
+ int N, OC_outer, OH, OW, IC_outer, FH, FW, sH, sW;
+ } conv_shape_t;
+ const conv_shape_t shapes[] = {
+ // An output extent is 0, so exit early.
+ {0, 1, 3, 3, 1, 3, 3, 1, 1}, // N = 0
+ {1, 0, 3, 3, 1, 3, 3, 1, 1}, // OC_outer = 0
+ {1, 1, 0, 3, 1, 3, 3, 1, 1}, // OH = 0
+ {1, 1, 3, 0, 1, 3, 3, 1, 1}, // OW = 0
+ // Zero reduction: non-accumulate must zero the output panel; accumulate
+ // must leave it intact.
+ {1, 1, 2, 2, 0, 3, 3, 1, 1}, // IC_outer = 0
+ {1, 1, 2, 2, 1, 0, 3, 1, 1}, // FH = 0
+ {1, 1, 2, 2, 1, 3, 0, 1, 1}, // FW = 0
+ // Regular cases.
+ {1, 1, 1, 1, 1, 1, 1, 1, 1},
+ {1, 1, 1, 15, 1, 1, 1, 1, 1},
+ {1, 1, 1, 16, 1, 1, 1, 1, 1},
+ {1, 1, 1, 17, 1, 1, 1, 1, 1},
+ {1, 2, 3, 5, 2, 3, 3, 1, 1},
+ {2, 2, 5, 5, 2, 3, 3, 1, 1},
+ {1, 1, 3, 3, 1, 1, 1, 1, 1},
+ {1, 1, 5, 5, 1, 3, 3, 2, 2},
+ {1, 1, 1, 33, 1, 3, 3, 1, 1},
+ };
+ for (int i = 0; i < (int)IREE_ARRAYSIZE(shapes); ++i) {
+ iree_uk_conv_nchwc_params_t params;
+ memcpy(¶ms, src_params, sizeof params);
+ params.cpu_data = iree_uk_test_cpu_data(test);
+ if (!(params.flags &
+ IREE_UK_FLAG_CONV_NCHWC_ALLOW_GENERIC_FALLBACK_TILE_FUNCTION)) {
+ if (!(iree_uk_conv_nchwc_info_p(¶ms) &
+ IREE_UK_FLAG_CONV_NCHWC_INFO_HAVE_ARCHITECTURE_SPECIFIC_TILE_FUNCTION)) {
+ IREE_UK_ASSERT(0 &&
+ "No architecture-specific tile function available "
+ "for this case or missing CPU features, which should "
+ "have been handled earlier.");
+ IREE_UK_TEST_FAIL(test);
+ }
+ }
+ conv_shape_t s = shapes[i];
+ params.N = s.N;
+ params.OC_outer = s.OC_outer;
+ params.OH = s.OH;
+ params.OW = s.OW;
+ params.IC_outer = s.IC_outer;
+ params.FH = s.FH;
+ params.FW = s.FW;
+ params.stride_h = s.sH;
+ params.stride_w = s.sW;
+ const iree_uk_uint32_t base_flags = params.flags;
+ for (int accumulate = 0; accumulate <= 1; ++accumulate) {
+ params.flags = base_flags;
+ if (accumulate) params.flags |= IREE_UK_FLAG_CONV_NCHWC_ACCUMULATE;
+ iree_uk_test_conv_nchwc_for_shape_params(test, ¶ms);
+ }
+ }
+}
+
+static void iree_uk_test_conv_nchwc(iree_uk_uint32_t flags, int k0, int c0,
+ const char* cpu_features) {
+ // Always allow the fallback in this test. The problem with trying to enforce
+ // that no fallback is accidentally used, is that it's not easy to tell when
+ // a fallback is legitimate. It depends on the build system used (CMake or
+ // Bazel) and within the CMake case, it depends on the native toolchain. The
+ // ground truth is given by the IREE_UK_BUILD_* defines, but they are
+ // currently an implementation detail within each arch/ subdirectory.
+ flags |= IREE_UK_FLAG_CONV_NCHWC_ALLOW_GENERIC_FALLBACK_TILE_FUNCTION;
+ char types_str[32];
+ iree_uk_conv_nchwc_type_t type = iree_uk_conv_nchwc_type(flags);
+ iree_uk_type_triple_str(types_str, sizeof types_str, type);
+ iree_uk_conv_nchwc_params_t params = {.flags = flags, .k0 = k0, .c0 = c0};
+ char test_label_str[256];
+ snprintf(test_label_str, sizeof test_label_str, "types:%s tile:%dx%d",
+ types_str, k0, c0);
+ iree_uk_test(test_label_str, iree_uk_test_conv_nchwc_for_tile_params, ¶ms,
+ cpu_features);
+}
+
+int main(int argc, char** argv) {
+ // Generic (no CPU features). Exercises the generic fallback tile.
+ iree_uk_test_conv_nchwc(IREE_UK_FLAG_CONV_NCHWC_TYPE_F32F32F32, 16, 16, "");
+ // A tile shape no arch dispatch table matches — exercises the generic
+ // fallback even when arch-specific tiles are compiled in.
+ iree_uk_test_conv_nchwc(IREE_UK_FLAG_CONV_NCHWC_TYPE_F32F32F32, 4, 4, "");
+
+#if defined(IREE_ARCH_X86_64)
+ iree_uk_test_conv_nchwc(IREE_UK_FLAG_CONV_NCHWC_TYPE_F32F32F32, 16, 16,
+ "avx512_base");
+#endif
+
+ return iree_uk_test_exit_status();
+}
