diff --git a/.gitignore b/.gitignore index ec6f4f79..30c1e1d5 100755 --- a/.gitignore +++ b/.gitignore @@ -23,3 +23,6 @@ id_ed25519.pub *.egg-info **/*.prj/** /outputs/ + +# Local agent tooling +.grok/ diff --git a/aie_kernels/aie2/conv2d.cc b/aie_kernels/aie2/conv2d.cc new file mode 100644 index 00000000..6ca72c32 --- /dev/null +++ b/aie_kernels/aie2/conv2d.cc @@ -0,0 +1,487 @@ +// SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +// 2D Convolution Kernel for AIE2 (NPU) +// Supports standard conv2d with configurable kernel_size, stride, padding + +#define NOCPP + +#include "../aie_kernel_utils.h" + +#include +#include +#include +#include + +extern "C" { + +/** + * 2D Convolution Kernel - AIE2 optimized + * Naive implementation for small kernels (3x3, 5x5) + * + * @param input - Input tensor [in_channels * in_height * in_width] + * @param weight - Weight tensor [out_channels * in_channels * kernel_height * kernel_width] + * @param output - Output tensor [out_channels * out_height * out_width] + * @param bias - Optional bias tensor [out_channels], can be NULL + * @param in_channels - Number of input channels + * @param in_height - Input height + * @param in_width - Input width + * @param out_channels - Number of output channels + * @param out_height - Output height + * @param out_width - Output width + * @param kernel_height - Kernel height + * @param kernel_width - Kernel width + * @param stride_height - Stride in height dimension + * @param stride_width - Stride in width dimension + * @param pad_height - Padding in height dimension + * @param pad_width - Padding in width dimension + */ +void conv2d_bf16_scalar(bfloat16 *input, + bfloat16 *weight, + bfloat16 *output, + bfloat16 *bias, + int in_channels, + int in_height, + int in_width, + int out_channels, + int out_height, + int out_width, + int kernel_height, + int kernel_width, + int stride_height, + int stride_width, + int pad_height, + int pad_width, + int groups, + int apply_bias) +{ + int channels_per_group = in_channels / groups; + int out_channels_per_group = out_channels / groups; + + for (int oc = 0; oc < out_channels; oc++) { + int group_id = oc / out_channels_per_group; + int oc_in_group = oc % out_channels_per_group; + + for (int oh = 0; oh < out_height; oh++) { + for (int ow = 0; ow < out_width; ow++) { + // Calculate input position + int ih_start = oh * stride_height - pad_height; + int iw_start = ow * stride_width - pad_width; + + bfloat16 acc = bfloat16(0.0f); + + // Sum over input channels in the group + for (int ic = 0; ic < channels_per_group; ic++) { + int ic_global = group_id * channels_per_group + ic; + + for (int kh = 0; kh < kernel_height; kh++) { + for (int kw = 0; kw < kernel_width; kw++) { + int ih = ih_start + kh * 1; // dilation = 1 for now + int iw = iw_start + kw * 1; + + // Check bounds (handle padding) + if (ih >= 0 && ih < in_height && iw >= 0 && iw < in_width) { + // NCHW flat: (ic_global * H + ih) * W + iw (N=1 layout) + int input_idx = (ic_global * in_height + ih) * in_width + iw; + int weight_idx = + ((oc * channels_per_group + ic) * kernel_height + kh) * kernel_width + kw; + + acc += input[input_idx] * weight[weight_idx]; + } + } + } + } + + // Packed bias: B_tile follows W_tile in the weight buffer. + if (apply_bias) { + int w_only = out_channels * channels_per_group * kernel_height * kernel_width; + acc += weight[w_only + oc]; + (void)bias; + } + + int output_idx = (oc * out_height + oh) * out_width + ow; + output[output_idx] = acc; + } + } + } +} + +/** + * 2D Convolution Kernel - Vectorized version for AIE2 + * + * Pointwise (k=1): vectorize over OW with aie::mac float accum when stride_w==1. + * k>1: scalar float path (OW-vector is unsafe for padded/H-strip RF widths). + */ +void conv2d_bf16_vector(bfloat16 *input, + bfloat16 *weight, + bfloat16 *output, + bfloat16 *bias, + int N, // batch size + int in_channels, + int in_height, + int in_width, + int out_channels, + int out_height, + int out_width, + int kernel_h, + int kernel_w, + int stride_h, + int stride_w, + int pad_h, + int pad_w, + int groups, + int apply_bias) +{ + constexpr int vec_factor = 8; + + event0(); + + int channels_per_group = in_channels / groups; + int out_channels_per_group = out_channels / groups; + int spatial_size = out_height * out_width; + const int w_only = out_channels * channels_per_group * kernel_h * kernel_w; + const aie::vector ones = aie::broadcast(bfloat16(1.0f)); + + for (int n = 0; n < N; n++) { + for (int oc = 0; oc < out_channels; oc++) { + int group_id = oc / out_channels_per_group; + int ic_start = group_id * channels_per_group; + + bfloat16 *__restrict out_ptr = output + ((n * out_channels + oc) * spatial_size); + const bfloat16 *__restrict w_oc = weight + oc * channels_per_group * kernel_h * kernel_w; + + for (int oh = 0; oh < out_height; oh++) { + int ih_base = oh * stride_h - pad_h; + int ow = 0; + + // Dense OW vector only for k=1; k>1 uses scalar float below. + if (stride_w == 1 && kernel_h == 1 && kernel_w == 1) { + for (; ow + vec_factor <= out_width; ow += vec_factor) { + aie::accum acc = aie::zeros(); + + if (apply_bias) { + acc = aie::mac(acc, aie::broadcast(weight[w_only + oc]), ones); + (void)bias; + } + + for (int ic = 0; ic < channels_per_group; ic++) { + int ic_global = ic_start + ic; + const bfloat16 *__restrict in_ch = + input + (n * in_channels + ic_global) * in_height * in_width; + + for (int kh = 0; kh < kernel_h; kh++) { + int ih = ih_base + kh; + if (ih < 0 || ih >= in_height) { + continue; + } + const bfloat16 *__restrict in_row = in_ch + ih * in_width; + + for (int kw = 0; kw < kernel_w; kw++) { + int iw0 = ow - pad_w + kw; + int iw_last = iw0 + vec_factor - 1; + aie::vector w_vec = aie::broadcast( + w_oc[(ic * kernel_h + kh) * kernel_w + kw]); + aie::vector in_vec; + // Do not write vector lanes via operator[] (not reliable on AIE). + // Interior aligned → load_v; else gather into aligned tmp then load_v. + if (iw0 >= 0 && iw_last < in_width && (iw0 & (vec_factor - 1)) == 0) { + in_vec = aie::load_v(in_row + iw0); + } else { + alignas(32) bfloat16 gather_tmp[vec_factor]; + for (int i = 0; i < vec_factor; i++) { + int iw = iw0 + i; + gather_tmp[i] = (iw >= 0 && iw < in_width) ? in_row[iw] : bfloat16(0.0f); + } + in_vec = aie::load_v(gather_tmp); + } + + acc = aie::mac(acc, in_vec, w_vec); + } + } + } + + aie::vector out_vec = acc.template to_vector(); + int out_off = oh * out_width + ow; + for (int i = 0; i < vec_factor; i++) { + out_ptr[out_off + i] = out_vec[i]; + } + } + } + + for (; ow < out_width; ow++) { + int ih_start = oh * stride_h - pad_h; + int iw_start = ow * stride_w - pad_w; + float acc = apply_bias ? float(weight[w_only + oc]) : 0.0f; + if (apply_bias) { + (void)bias; + } + + for (int ic = 0; ic < channels_per_group; ic++) { + int ic_global = ic_start + ic; + for (int kh = 0; kh < kernel_h; kh++) { + for (int kw = 0; kw < kernel_w; kw++) { + int ih = ih_start + kh; + int iw = iw_start + kw; + if (ih >= 0 && ih < in_height && iw >= 0 && iw < in_width) { + int input_idx = ((n * in_channels + ic_global) * in_height + ih) * in_width + iw; + int weight_idx = ((oc * channels_per_group + ic) * kernel_h + kh) * kernel_w + kw; + acc += float(input[input_idx]) * float(weight[weight_idx]); + } + } + } + } + out_ptr[oh * out_width + ow] = static_cast(acc); + } + } + } + } + + event1(); +} + +/** + * Depthwise Convolution Kernel - AIE2 (vec width 8) + * OW-dense pipeline parity with aie2p depthwise. + */ +void depthwise_conv2d_bf16_vector(bfloat16 *input, + bfloat16 *weight, + bfloat16 *output, + bfloat16 *bias, + int N, + int channels, + int in_height, + int in_width, + int out_height, + int out_width, + int kernel_h, + int kernel_w, + int stride_h, + int stride_w, + int pad_h, + int pad_w, + int apply_bias) +{ + constexpr int vec_factor = 8; + + event0(); + + int spatial_size = out_height * out_width; + const int w_only = channels * kernel_h * kernel_w; + const aie::vector ones = aie::broadcast(bfloat16(1.0f)); + + for (int n = 0; n < N; n++) { + for (int c = 0; c < channels; c++) { + bfloat16 *__restrict out_ptr = output + (n * channels + c) * spatial_size; + const bfloat16 *__restrict in_ch = input + (n * channels + c) * in_height * in_width; + const bfloat16 *__restrict w_c = weight + c * kernel_h * kernel_w; + + for (int oh = 0; oh < out_height; oh++) { + int ih_base = oh * stride_h - pad_h; + int ow = 0; + + if (stride_w == 1) { + for (; ow + vec_factor <= out_width; ow += vec_factor) { + aie::accum acc = aie::zeros(); + + if (apply_bias) { + acc = aie::mac(acc, aie::broadcast(weight[w_only + c]), ones); + (void)bias; + } + + for (int kh = 0; kh < kernel_h; kh++) { + int ih = ih_base + kh; + if (ih < 0 || ih >= in_height) { + continue; + } + const bfloat16 *__restrict in_row = in_ch + ih * in_width; + + for (int kw = 0; kw < kernel_w; kw++) { + int iw0 = ow - pad_w + kw; + aie::vector w_vec = + aie::broadcast(w_c[kh * kernel_w + kw]); + aie::vector in_vec; + + int iw_last = iw0 + vec_factor - 1; + // Avoid vector lane operator[] writes; gather via aligned tmp. + if (iw0 >= 0 && iw_last < in_width && (iw0 & (vec_factor - 1)) == 0) { + in_vec = aie::load_v(in_row + iw0); + } else { + alignas(32) bfloat16 gather_tmp[vec_factor]; + for (int i = 0; i < vec_factor; i++) { + int iw = iw0 + i; + gather_tmp[i] = (iw >= 0 && iw < in_width) ? in_row[iw] : bfloat16(0.0f); + } + in_vec = aie::load_v(gather_tmp); + } + + acc = aie::mac(acc, in_vec, w_vec); + } + } + + aie::vector out_vec = acc.template to_vector(); + int out_off = oh * out_width + ow; + for (int i = 0; i < vec_factor; i++) { + out_ptr[out_off + i] = out_vec[i]; + } + } + } + + for (; ow < out_width; ow++) { + int ih_start = oh * stride_h - pad_h; + int iw_start = ow * stride_w - pad_w; + float acc = apply_bias ? float(weight[w_only + c]) : 0.0f; + if (apply_bias) { + (void)bias; + } + + for (int kh = 0; kh < kernel_h; kh++) { + for (int kw = 0; kw < kernel_w; kw++) { + int ih = ih_start + kh; + int iw = iw_start + kw; + if (ih >= 0 && ih < in_height && iw >= 0 && iw < in_width) { + int input_idx = ((n * channels + c) * in_height + ih) * in_width + iw; + acc += float(input[input_idx]) * float(w_c[kh * kernel_w + kw]); + } + } + } + out_ptr[oh * out_width + ow] = static_cast(acc); + } + } + } + } + + event1(); +} + +/** + * Pointwise (1x1) Convolution — AIE2 dense path + */ + * + * NCHW: contiguous loads over H*W channel planes with broadcast weight[oc, ic] + * (see aie2p pointwise for layout rationale). Native vector width 8 bf16. + * + * @param input - Input tensor [N, in_channels, H, W] + * @param weight - Weight tensor [out_channels, in_channels] (+ packed bias) + * @param output - Output tensor [N, out_channels, H, W] + * @param bias - Unused when bias is packed after weights + */ +void pointwise_conv2d_bf16_vector(bfloat16 *input, + bfloat16 *weight, + bfloat16 *output, + bfloat16 *bias, + int N, + int in_channels, + int out_channels, + int height, + int width, + int apply_bias) + { + constexpr int vec_factor = 8; + constexpr int oc_tile = 4; + + event0(); + + const int spatial_size = height * width; + const int w_only = out_channels * in_channels; + const aie::vector ones = aie::broadcast(bfloat16(1.0f)); + + for (int n = 0; n < N; n++) { + bfloat16 *__restrict in_n = input + n * in_channels * spatial_size; + bfloat16 *__restrict out_n = output + n * out_channels * spatial_size; + + int oc = 0; + for (; oc + oc_tile <= out_channels; oc += oc_tile) { + const bfloat16 *__restrict w0 = weight + (oc + 0) * in_channels; + const bfloat16 *__restrict w1 = weight + (oc + 1) * in_channels; + const bfloat16 *__restrict w2 = weight + (oc + 2) * in_channels; + const bfloat16 *__restrict w3 = weight + (oc + 3) * in_channels; + bfloat16 *__restrict o0 = out_n + (oc + 0) * spatial_size; + bfloat16 *__restrict o1 = out_n + (oc + 1) * spatial_size; + bfloat16 *__restrict o2 = out_n + (oc + 2) * spatial_size; + bfloat16 *__restrict o3 = out_n + (oc + 3) * spatial_size; + + int sp = 0; + for (; sp + vec_factor <= spatial_size; sp += vec_factor) { + aie::accum a0 = aie::zeros(); + aie::accum a1 = aie::zeros(); + aie::accum a2 = aie::zeros(); + aie::accum a3 = aie::zeros(); + + if (apply_bias) { + a0 = aie::mac(a0, aie::broadcast(weight[w_only + oc + 0]), ones); + a1 = aie::mac(a1, aie::broadcast(weight[w_only + oc + 1]), ones); + a2 = aie::mac(a2, aie::broadcast(weight[w_only + oc + 2]), ones); + a3 = aie::mac(a3, aie::broadcast(weight[w_only + oc + 3]), ones); + (void)bias; + } + + for (int ic = 0; ic < in_channels; ic++) { + aie::vector in_vec = aie::load_v(in_n + ic * spatial_size + sp); + a0 = aie::mac(a0, in_vec, aie::broadcast(w0[ic])); + a1 = aie::mac(a1, in_vec, aie::broadcast(w1[ic])); + a2 = aie::mac(a2, in_vec, aie::broadcast(w2[ic])); + a3 = aie::mac(a3, in_vec, aie::broadcast(w3[ic])); + } + + aie::store_v(o0 + sp, a0.template to_vector()); + aie::store_v(o1 + sp, a1.template to_vector()); + aie::store_v(o2 + sp, a2.template to_vector()); + aie::store_v(o3 + sp, a3.template to_vector()); + } + + for (; sp < spatial_size; sp++) { + float f0 = 0.0f, f1 = 0.0f, f2 = 0.0f, f3 = 0.0f; + if (apply_bias) { + f0 = weight[w_only + oc + 0]; + f1 = weight[w_only + oc + 1]; + f2 = weight[w_only + oc + 2]; + f3 = weight[w_only + oc + 3]; + (void)bias; + } + for (int ic = 0; ic < in_channels; ic++) { + float x = in_n[ic * spatial_size + sp]; + f0 += x * float(w0[ic]); + f1 += x * float(w1[ic]); + f2 += x * float(w2[ic]); + f3 += x * float(w3[ic]); + } + o0[sp] = static_cast(f0); + o1[sp] = static_cast(f1); + o2[sp] = static_cast(f2); + o3[sp] = static_cast(f3); + } + } + + for (; oc < out_channels; oc++) { + const bfloat16 *__restrict w_row = weight + oc * in_channels; + bfloat16 *__restrict out_ptr = out_n + oc * spatial_size; + + int sp = 0; + for (; sp + vec_factor <= spatial_size; sp += vec_factor) { + aie::accum acc = aie::zeros(); + if (apply_bias) { + acc = aie::mac(acc, aie::broadcast(weight[w_only + oc]), ones); + (void)bias; + } + for (int ic = 0; ic < in_channels; ic++) { + aie::vector in_vec = aie::load_v(in_n + ic * spatial_size + sp); + acc = aie::mac(acc, in_vec, aie::broadcast(w_row[ic])); + } + aie::store_v(out_ptr + sp, acc.template to_vector()); + } + for (; sp < spatial_size; sp++) { + float f = apply_bias ? float(weight[w_only + oc]) : 0.0f; + if (apply_bias) { + (void)bias; + } + for (int ic = 0; ic < in_channels; ic++) { + f += float(in_n[ic * spatial_size + sp]) * float(w_row[ic]); + } + out_ptr[sp] = static_cast(f); + } + } + } + + event1(); + } + } // extern "C" diff --git a/aie_kernels/aie2p/conv2d.cc b/aie_kernels/aie2p/conv2d.cc new file mode 100644 index 00000000..51a14834 --- /dev/null +++ b/aie_kernels/aie2p/conv2d.cc @@ -0,0 +1,493 @@ +// SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +// 2D Convolution Kernel for AIE2P (NPU2) + +#define NOCPP + +#include "../aie_kernel_utils.h" + +#include +#include +#include +#include + +extern "C" { + +/** + * 2D Convolution Kernel - AIE2P optimized + * Vector factor 16 (AIE2P). + * + * @param input - Input tensor [N, in_channels, in_height, in_width] (flattened) + * @param weight - Weight tensor [out_channels, in_channels, kernel_height, kernel_width] + * @param output - Output tensor [N, out_channels, out_height, out_width] (flattened) + * @param bias - Optional bias tensor [out_channels] + */ +void conv2d_bf16_scalar(bfloat16 *input, + bfloat16 *weight, + bfloat16 *output, + bfloat16 *bias, + int N, // batch size + int in_channels, + int in_height, + int in_width, + int out_channels, + int out_height, + int out_width, + int kernel_h, + int kernel_w, + int stride_h, + int stride_w, + int pad_h, + int pad_w, + int groups, + int apply_bias) +{ + int channels_per_group = in_channels / groups; + int out_channels_per_group = out_channels / groups; + + for (int n = 0; n < N; n++) { + for (int oc = 0; oc < out_channels; oc++) { + int group_id = oc / out_channels_per_group; + int ic_start = group_id * channels_per_group; + + for (int oh = 0; oh < out_height; oh++) { + for (int ow = 0; ow < out_width; ow++) { + int ih_start = oh * stride_h - pad_h; + int iw_start = ow * stride_w - pad_w; + + bfloat16 acc = bfloat16(0.0f); + + for (int ic = 0; ic < channels_per_group; ic++) { + int ic_global = ic_start + ic; + + for (int kh = 0; kh < kernel_h; kh++) { + for (int kw = 0; kw < kernel_w; kw++) { + int ih = ih_start + kh; + int iw = iw_start + kw; + + if (ih >= 0 && ih < in_height && iw >= 0 && iw < in_width) { + int input_idx = ((n * in_channels + ic_global) * in_height + ih) * in_width + iw; + int weight_idx = ((oc * channels_per_group + ic) * kernel_h + kh) * kernel_w + kw; + + acc += input[input_idx] * weight[weight_idx]; + } + } + } + } + + if (apply_bias) { + // Packed bias: B_tile follows W_tile in the weight buffer. + int w_only = out_channels * channels_per_group * kernel_h * kernel_w; + acc += weight[w_only + oc]; + (void)bias; + } + + int out_idx = ((n * out_channels + oc) * out_height + oh) * out_width + ow; + output[out_idx] = acc; + } + } + } + } +} + +/** + * 2D Convolution Kernel - Vectorized version for AIE2P + * + * Pointwise (k=1): vectorize over OW with aie::mac float accum when stride_w==1. + * k>1 / non-unit stride_w: scalar float path. + * + * @param input - Input tensor [N, in_channels, in_height, in_width] (flattened) + * @param weight - Weight tensor [out_channels, in_channels, kernel_height, kernel_width] + * @param output - Output tensor [N, out_channels, out_height, out_width] (flattened) + * @param bias - Optional bias tensor [out_channels] + */ +void conv2d_bf16_vector(bfloat16 *input, + bfloat16 *weight, + bfloat16 *output, + bfloat16 *bias, + int N, // batch size + int in_channels, + int in_height, + int in_width, + int out_channels, + int out_height, + int out_width, + int kernel_h, + int kernel_w, + int stride_h, + int stride_w, + int pad_h, + int pad_w, + int groups, + int apply_bias) +{ + constexpr int vec_factor = 16; + + event0(); + + int channels_per_group = in_channels / groups; + int out_channels_per_group = out_channels / groups; + int spatial_size = out_height * out_width; + const int w_only = out_channels * channels_per_group * kernel_h * kernel_w; + const aie::vector ones = aie::broadcast(bfloat16(1.0f)); + + for (int n = 0; n < N; n++) { + for (int oc = 0; oc < out_channels; oc++) { + int group_id = oc / out_channels_per_group; + int ic_start = group_id * channels_per_group; + + bfloat16 *__restrict out_ptr = output + (n * out_channels + oc) * spatial_size; + const bfloat16 *__restrict w_oc = weight + oc * channels_per_group * kernel_h * kernel_w; + + for (int oh = 0; oh < out_height; oh++) { + int ih_base = oh * stride_h - pad_h; + int ow = 0; + + // Dense OW tiles: unit stride in W → contiguous input window. + // Dense OW vector only for k=1; k>1 uses scalar float below. + if (stride_w == 1 && kernel_h == 1 && kernel_w == 1) { + for (; ow + vec_factor <= out_width; ow += vec_factor) { + aie::accum acc = aie::zeros(); + + if (apply_bias) { + acc = aie::mac(acc, aie::broadcast(weight[w_only + oc]), ones); + (void)bias; + } + + for (int ic = 0; ic < channels_per_group; ic++) { + int ic_global = ic_start + ic; + const bfloat16 *__restrict in_ch = + input + (n * in_channels + ic_global) * in_height * in_width; + + for (int kh = 0; kh < kernel_h; kh++) { + int ih = ih_base + kh; + if (ih < 0 || ih >= in_height) { + continue; + } + const bfloat16 *__restrict in_row = in_ch + ih * in_width; + + for (int kw = 0; kw < kernel_w; kw++) { + int iw0 = ow - pad_w + kw; + int iw_last = iw0 + vec_factor - 1; + aie::vector w_vec = aie::broadcast( + w_oc[(ic * kernel_h + kh) * kernel_w + kw]); + aie::vector in_vec; + // Do not write vector lanes via operator[] (not reliable on AIE). + // Interior aligned → load_v; else gather into aligned tmp then load_v. + if (iw0 >= 0 && iw_last < in_width && (iw0 & (vec_factor - 1)) == 0) { + in_vec = aie::load_v(in_row + iw0); + } else { + alignas(32) bfloat16 gather_tmp[vec_factor]; + for (int i = 0; i < vec_factor; i++) { + int iw = iw0 + i; + gather_tmp[i] = (iw >= 0 && iw < in_width) ? in_row[iw] : bfloat16(0.0f); + } + in_vec = aie::load_v(gather_tmp); + } + + acc = aie::mac(acc, in_vec, w_vec); + } + } + } + + aie::vector out_vec = acc.template to_vector(); + int out_off = oh * out_width + ow; + for (int i = 0; i < vec_factor; i++) { + out_ptr[out_off + i] = out_vec[i]; + } + } + } + + // OW tail and non-unit stride_w: scalar float accum. + for (; ow < out_width; ow++) { + int ih_start = oh * stride_h - pad_h; + int iw_start = ow * stride_w - pad_w; + float acc = apply_bias ? float(weight[w_only + oc]) : 0.0f; + if (apply_bias) { + (void)bias; + } + + for (int ic = 0; ic < channels_per_group; ic++) { + int ic_global = ic_start + ic; + for (int kh = 0; kh < kernel_h; kh++) { + for (int kw = 0; kw < kernel_w; kw++) { + int ih = ih_start + kh; + int iw = iw_start + kw; + if (ih >= 0 && ih < in_height && iw >= 0 && iw < in_width) { + int input_idx = ((n * in_channels + ic_global) * in_height + ih) * in_width + iw; + int weight_idx = ((oc * channels_per_group + ic) * kernel_h + kh) * kernel_w + kw; + acc += float(input[input_idx]) * float(weight[weight_idx]); + } + } + } + } + out_ptr[oh * out_width + ow] = static_cast(acc); + } + } + } + } + + event1(); +} + +/** + * Depthwise Convolution Kernel - AIE2P optimized + * + * Same OW-dense pipeline as standard conv (stride_w==1 contiguous W loads), + * one input channel per output channel. + * + * @param input - Input tensor [N, channels, in_height, in_width] + * @param weight - Weight tensor [channels, kernel_h, kernel_w] + * @param output - Output tensor [N, channels, out_height, out_width] + * @param bias - Optional bias tensor [channels] + */ +void depthwise_conv2d_bf16_vector(bfloat16 *input, + bfloat16 *weight, + bfloat16 *output, + bfloat16 *bias, + int N, + int channels, + int in_height, + int in_width, + int out_height, + int out_width, + int kernel_h, + int kernel_w, + int stride_h, + int stride_w, + int pad_h, + int pad_w, + int apply_bias) +{ + constexpr int vec_factor = 16; + + event0(); + + int spatial_size = out_height * out_width; + const int w_only = channels * kernel_h * kernel_w; + const aie::vector ones = aie::broadcast(bfloat16(1.0f)); + + for (int n = 0; n < N; n++) { + for (int c = 0; c < channels; c++) { + bfloat16 *__restrict out_ptr = output + (n * channels + c) * spatial_size; + const bfloat16 *__restrict in_ch = input + (n * channels + c) * in_height * in_width; + const bfloat16 *__restrict w_c = weight + c * kernel_h * kernel_w; + + for (int oh = 0; oh < out_height; oh++) { + int ih_base = oh * stride_h - pad_h; + int ow = 0; + + if (stride_w == 1) { + for (; ow + vec_factor <= out_width; ow += vec_factor) { + aie::accum acc = aie::zeros(); + + if (apply_bias) { + acc = aie::mac(acc, aie::broadcast(weight[w_only + c]), ones); + (void)bias; + } + + for (int kh = 0; kh < kernel_h; kh++) { + int ih = ih_base + kh; + if (ih < 0 || ih >= in_height) { + continue; + } + const bfloat16 *__restrict in_row = in_ch + ih * in_width; + + for (int kw = 0; kw < kernel_w; kw++) { + int iw0 = ow - pad_w + kw; + aie::vector w_vec = + aie::broadcast(w_c[kh * kernel_w + kw]); + aie::vector in_vec; + + int iw_last = iw0 + vec_factor - 1; + // Avoid vector lane operator[] writes; gather via aligned tmp. + if (iw0 >= 0 && iw_last < in_width && (iw0 & (vec_factor - 1)) == 0) { + in_vec = aie::load_v(in_row + iw0); + } else { + alignas(32) bfloat16 gather_tmp[vec_factor]; + for (int i = 0; i < vec_factor; i++) { + int iw = iw0 + i; + gather_tmp[i] = (iw >= 0 && iw < in_width) ? in_row[iw] : bfloat16(0.0f); + } + in_vec = aie::load_v(gather_tmp); + } + + acc = aie::mac(acc, in_vec, w_vec); + } + } + + aie::vector out_vec = acc.template to_vector(); + int out_off = oh * out_width + ow; + for (int i = 0; i < vec_factor; i++) { + out_ptr[out_off + i] = out_vec[i]; + } + } + } + + for (; ow < out_width; ow++) { + int ih_start = oh * stride_h - pad_h; + int iw_start = ow * stride_w - pad_w; + float acc = apply_bias ? float(weight[w_only + c]) : 0.0f; + if (apply_bias) { + (void)bias; + } + + for (int kh = 0; kh < kernel_h; kh++) { + for (int kw = 0; kw < kernel_w; kw++) { + int ih = ih_start + kh; + int iw = iw_start + kw; + if (ih >= 0 && ih < in_height && iw >= 0 && iw < in_width) { + int input_idx = ((n * channels + c) * in_height + ih) * in_width + iw; + acc += float(input[input_idx]) * float(w_c[kh * kernel_w + kw]); + } + } + } + out_ptr[oh * out_width + ow] = static_cast(acc); + } + } + } + } + + event1(); +} + +/** + * Pointwise (1x1) Convolution — AIE2P dense path + * + * NCHW layout: channel planes are contiguous in H*W, IC is strided by spatial. + * Dense strategy (no host re-layout): treat each OC as a channel-plane axpy + * chain — contiguous `aie::load_v` over spatial, broadcast weight[oc, ic], + * float accum via `aie::mac`, store vector. Tile a few OCs so one input + * vector is reused (outer-product style), which is the NCHW-friendly dense + * pipeline. + * + * @param input - Input tensor [N, in_channels, H, W] + * @param weight - Weight tensor [out_channels, in_channels] (+ packed bias) + * @param output - Output tensor [N, out_channels, H, W] + * @param bias - Unused when bias is packed after weights + */ +void pointwise_conv2d_bf16_vector(bfloat16 *input, + bfloat16 *weight, + bfloat16 *output, + bfloat16 *bias, + int N, + int in_channels, + int out_channels, + int height, + int width, + int apply_bias) +{ + // AIE2P: 16-lane bf16 vectors for contiguous spatial planes. + constexpr int vec_factor = 16; + // How many output channels share one input spatial load. + constexpr int oc_tile = 4; + + event0(); + + const int spatial_size = height * width; + const int w_only = out_channels * in_channels; + const aie::vector ones = aie::broadcast(bfloat16(1.0f)); + + for (int n = 0; n < N; n++) { + bfloat16 *__restrict in_n = input + n * in_channels * spatial_size; + bfloat16 *__restrict out_n = output + n * out_channels * spatial_size; + + int oc = 0; + for (; oc + oc_tile <= out_channels; oc += oc_tile) { + const bfloat16 *__restrict w0 = weight + (oc + 0) * in_channels; + const bfloat16 *__restrict w1 = weight + (oc + 1) * in_channels; + const bfloat16 *__restrict w2 = weight + (oc + 2) * in_channels; + const bfloat16 *__restrict w3 = weight + (oc + 3) * in_channels; + bfloat16 *__restrict o0 = out_n + (oc + 0) * spatial_size; + bfloat16 *__restrict o1 = out_n + (oc + 1) * spatial_size; + bfloat16 *__restrict o2 = out_n + (oc + 2) * spatial_size; + bfloat16 *__restrict o3 = out_n + (oc + 3) * spatial_size; + + int sp = 0; + for (; sp + vec_factor <= spatial_size; sp += vec_factor) { + aie::accum a0 = aie::zeros(); + aie::accum a1 = aie::zeros(); + aie::accum a2 = aie::zeros(); + aie::accum a3 = aie::zeros(); + + if (apply_bias) { + a0 = aie::mac(a0, aie::broadcast(weight[w_only + oc + 0]), ones); + a1 = aie::mac(a1, aie::broadcast(weight[w_only + oc + 1]), ones); + a2 = aie::mac(a2, aie::broadcast(weight[w_only + oc + 2]), ones); + a3 = aie::mac(a3, aie::broadcast(weight[w_only + oc + 3]), ones); + (void)bias; + } + + for (int ic = 0; ic < in_channels; ic++) { + aie::vector in_vec = aie::load_v(in_n + ic * spatial_size + sp); + a0 = aie::mac(a0, in_vec, aie::broadcast(w0[ic])); + a1 = aie::mac(a1, in_vec, aie::broadcast(w1[ic])); + a2 = aie::mac(a2, in_vec, aie::broadcast(w2[ic])); + a3 = aie::mac(a3, in_vec, aie::broadcast(w3[ic])); + } + + aie::store_v(o0 + sp, a0.template to_vector()); + aie::store_v(o1 + sp, a1.template to_vector()); + aie::store_v(o2 + sp, a2.template to_vector()); + aie::store_v(o3 + sp, a3.template to_vector()); + } + + // Spatial tail (H*W not multiple of vec_factor): scalar float accum. + for (; sp < spatial_size; sp++) { + float f0 = 0.0f, f1 = 0.0f, f2 = 0.0f, f3 = 0.0f; + if (apply_bias) { + f0 = weight[w_only + oc + 0]; + f1 = weight[w_only + oc + 1]; + f2 = weight[w_only + oc + 2]; + f3 = weight[w_only + oc + 3]; + (void)bias; + } + for (int ic = 0; ic < in_channels; ic++) { + float x = in_n[ic * spatial_size + sp]; + f0 += x * float(w0[ic]); + f1 += x * float(w1[ic]); + f2 += x * float(w2[ic]); + f3 += x * float(w3[ic]); + } + o0[sp] = static_cast(f0); + o1[sp] = static_cast(f1); + o2[sp] = static_cast(f2); + o3[sp] = static_cast(f3); + } + } + + // Remainder OC (not multiple of oc_tile). + for (; oc < out_channels; oc++) { + const bfloat16 *__restrict w_row = weight + oc * in_channels; + bfloat16 *__restrict out_ptr = out_n + oc * spatial_size; + + int sp = 0; + for (; sp + vec_factor <= spatial_size; sp += vec_factor) { + aie::accum acc = aie::zeros(); + if (apply_bias) { + acc = aie::mac(acc, aie::broadcast(weight[w_only + oc]), ones); + (void)bias; + } + for (int ic = 0; ic < in_channels; ic++) { + aie::vector in_vec = aie::load_v(in_n + ic * spatial_size + sp); + acc = aie::mac(acc, in_vec, aie::broadcast(w_row[ic])); + } + aie::store_v(out_ptr + sp, acc.template to_vector()); + } + for (; sp < spatial_size; sp++) { + float f = apply_bias ? float(weight[w_only + oc]) : 0.0f; + if (apply_bias) { + (void)bias; + } + for (int ic = 0; ic < in_channels; ic++) { + f += float(in_n[ic * spatial_size + sp]) * float(w_row[ic]); + } + out_ptr[sp] = static_cast(f); + } + } + } + + event1(); +} +} // extern "C" diff --git a/iron/common/__init__.py b/iron/common/__init__.py index cb2ff31b..500ebfe8 100644 --- a/iron/common/__init__.py +++ b/iron/common/__init__.py @@ -5,6 +5,7 @@ from .base import ( AIEOperatorBase, + AIEOperatorConstraintError, MLIROperator, CompositeOperator, AIERuntimeArgSpec, @@ -16,6 +17,8 @@ KernelArchiveArtifact, SourceArtifact, PythonGeneratedMLIRArtifact, + XclbinArtifact, + InstsBinArtifact, DesignGenerator, ) from .layout import Stride, TiledStride, TiledStridedLayout, tiled_2d diff --git a/iron/common/base.py b/iron/common/base.py index 701e90df..f2e4c39c 100644 --- a/iron/common/base.py +++ b/iron/common/base.py @@ -216,3 +216,15 @@ def __post_init__(self) -> None: raise ValueError( f"Invalid direction {self.direction!r}: must be one of 'in', 'out', 'inout'" ) + + +class AIEOperatorConstraintError(RuntimeError): + """Raised by AIE operators when runtime inputs violate constructor-time constraints + (e.g., shape, dtype, channel count, or spatial dimensions that were baked into the + compiled kernel at operator construction time). + + This allows clean separation between construction-time specialization and + runtime validation without using generic exceptions. + """ + + pass diff --git a/iron/operators/__init__.py b/iron/operators/__init__.py index 6d62e215..d81218ac 100644 --- a/iron/operators/__init__.py +++ b/iron/operators/__init__.py @@ -3,6 +3,7 @@ from .elementwise_add.op import ElementwiseAdd from .elementwise_mul.op import ElementwiseMul +from .conv2d.op import AIEConv2d from .gemm.op import GEMM from .gemv.op import GEMV from .mha.op import MHA diff --git a/iron/operators/conv2d/cpu_test.py b/iron/operators/conv2d/cpu_test.py new file mode 100644 index 00000000..5a26e39b --- /dev/null +++ b/iron/operators/conv2d/cpu_test.py @@ -0,0 +1,362 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Pure-CPU tests for AIEConv2d reference (no XRT/NPU).""" + +import pytest + +import torch +import torch.nn.functional as F + +from .reference import ( + generate_golden_reference, + conv2d_cpu, + calculate_output_dim, +) +from .test import get_params + + +@pytest.mark.parametrize( + "dummy", + [pytest.param(None, id="reference_cpu_only")], +) +def test_conv2d_reference_cpu_only(dummy): + """CPU-only checks of generate_golden_reference / conv2d_cpu / dims.""" + # Broad representative cases exercising all important golden + dim paths. + # All cases satisfy F.conv2d validity (spatials after pad >= kernel). + test_cases = [ + # (bs, ic, h, w, oc, k, s, p, g, use_bias) + (1, 3, 32, 32, 16, 3, 1, 1, 1, True), # basic bias (regular style) + (1, 3, 32, 32, 16, 3, 1, 1, 1, False), # basic nobias + (1, 16, 32, 32, 16, 3, 1, 1, 16, True), # depthwise +bias + (1, 16, 32, 32, 16, 3, 1, 1, 16, False), # depthwise nobias + (2, 32, 16, 16, 64, 1, 1, 0, 1, True), # pointwise + batch>1 + (1, 16, 32, 32, 32, 3, 2, 1, 1, True), # strided + pad + (1, 16, 32, 32, 32, 3, 2, 0, 1, True), # strided no pad + (1, 8, 8, 8, 16, 3, 1, 2, 2, True), # groups=2 + overhang pad + (1, 4, 7, 9, 8, 3, 1, 1, 2, False), # groups + small + nobias + (4, 4, 8, 8, 8, 1, 1, 0, 1, True), # batch + pointwise no pad + ] + + for bs, ic, h, w, oc, k, s, p, g, ub in test_cases: + golden = generate_golden_reference( + batch_size=bs, + in_channels=ic, + in_height=h, + in_width=w, + out_channels=oc, + kernel_size=k, + stride=s, + padding=p, + groups=g, + use_bias=ub, + seed=42 + hash((bs, ic, h, w, oc, k, s, p, g, ub)) % 10000, + ) + + # Direct authoritative ground truth + direct = F.conv2d( + golden["input"], + golden["weight"], + golden["bias"], + stride=s, + padding=p, + groups=g, + ) + + # Golden must match F.conv2d exactly (same contract as conv2d_cpu) + assert torch.equal( + golden["output"], direct + ), f"ref mismatch for case {(bs,ic,h,w,oc,k,s,p,g,ub)}" + + # Exercise conv2d_cpu wrapper itself (the one wrapped by golden) + cpu_out = conv2d_cpu( + golden["input"], golden["weight"], golden["bias"], s, p, 1, g + ) + assert torch.equal(cpu_out, golden["output"]) + + # Exercise calculate_output_dim (used by get_params for divis + naming) + calc_h = calculate_output_dim(h, k, s, p, 1) + calc_w = calculate_output_dim(w, k, s, p, 1) + assert calc_h == direct.shape[2] + assert calc_w == direct.shape[3] + + # Also match operator's internal formula (for cross-guard) + op_h = (h + 2 * p - k) // s + 1 + op_w = (w + 2 * p - k) // s + 1 + assert op_h == calc_h and op_w == calc_w + + # Live sanity: get_params / all_params must be healthy at collection time + all_p = get_params() + assert len(all_p) > 20, "get_params produced too few cases" + non_ext = [ + p + for p in all_p + if not any( + getattr(m, "name", None) == "extensive" for m in getattr(p, "marks", []) + ) + ] + assert len(non_ext) >= 1, "No regular (non-extensive) cases in matrix" + # The first regular must be unmarked + first_reg_marks = getattr(non_ext[0], "marks", []) + assert not any( + getattr(m, "name", None) == "extensive" for m in first_reg_marks + ), "First regular case unexpectedly marked extensive" + + print( + "\nConv2D pure CPU reference test: all cases PASS (exact matches + dim checks)." + ) + print(f" all_params count: {len(all_p)} (regular + extensive matrix healthy)") + + +# Explicit CPU_REFERENCE_CASES using production-grade pytest.param with stable ids. +# These mirror (and are a superset of) the families exercised by get_params and +# the forward tests. IDs are human-readable and safe for CSV/metrics reporting. +CPU_REFERENCE_CASES = [ + # Core + bias variants (matches regular matrix spirit) + pytest.param(1, 3, 32, 32, 16, 3, 1, 1, 1, True, 42, id="cpu_basic_bias"), + pytest.param(1, 3, 32, 32, 16, 3, 1, 1, 1, False, 42, id="cpu_basic_nobias"), + # Depthwise + pytest.param(1, 16, 32, 32, 16, 3, 1, 1, 16, True, 123, id="cpu_depthwise_bias"), + pytest.param(1, 16, 32, 32, 16, 3, 1, 1, 16, False, 123, id="cpu_depthwise_nobias"), + # Pointwise + pytest.param(1, 32, 32, 32, 64, 1, 1, 0, 1, True, 7, id="cpu_pointwise_bias"), + pytest.param(1, 32, 32, 32, 64, 1, 1, 0, 1, False, 7, id="cpu_pointwise_nobias"), + # Strided cases (p=0 and p=1) + pytest.param(1, 16, 32, 32, 32, 3, 2, 1, 1, True, 99, id="cpu_strided_p1"), + pytest.param(1, 16, 32, 32, 32, 3, 2, 0, 1, True, 99, id="cpu_strided_p0"), + # Grouped + pytest.param(1, 8, 16, 16, 16, 3, 1, 2, 2, True, 2026, id="cpu_groups2"), + pytest.param(1, 4, 16, 16, 8, 3, 1, 1, 2, True, 11, id="cpu_groups2_small"), + # batch > 1 (exercises generate path used by forward batch-2 test) + pytest.param(2, 3, 32, 32, 16, 3, 1, 1, 1, True, 55, id="cpu_batch2"), + pytest.param(3, 16, 16, 16, 16, 3, 1, 1, 16, False, 88, id="cpu_depthwise_batch3"), + # Different spatial + seed for reproducibility cross-check + pytest.param(1, 3, 64, 64, 16, 3, 1, 1, 1, True, 0, id="cpu_large_spatial"), +] + + +@pytest.mark.parametrize( + "batch,in_ch,h,w,out_ch,k,s,p,g,use_bias,seed", + CPU_REFERENCE_CASES, +) +def test_conv2d_cpu_reference_only( + batch, in_ch, h, w, out_ch, k, s, p, g, use_bias, seed +): + """Golden output matches direct conv2d_cpu on the same tensors.""" + # Via the golden path (what HW tests actually use) + golden = generate_golden_reference( + batch_size=batch, + in_channels=in_ch, + in_height=h, + in_width=w, + out_channels=out_ch, + kernel_size=k, + stride=s, + padding=p, + groups=g, + use_bias=use_bias, + dtype=torch.bfloat16, + seed=seed, + ) + via_golden = golden["output"] + + # Direct call to the CPU reference (thin F.conv2d wrapper) + direct = conv2d_cpu( + input=golden["input"], + weight=golden["weight"], + bias=golden["bias"], + stride=s, + padding=p, + dilation=1, + groups=g, + ) + + # Must be identical (same seed + same deterministic path through conv2d_cpu) + assert ( + direct.shape == via_golden.shape + ), f"Shape mismatch direct vs golden: {direct.shape} vs {via_golden.shape}" + assert direct.dtype == via_golden.dtype == torch.bfloat16 + + # Exact match expected (identical computation, no AIE involved) + assert torch.equal(direct, via_golden), ( + "conv2d_cpu direct result does not bitwise match golden['output'] " + "(the value passed to run_test / forward). This breaks the reference contract." + ) + + # Sanity: config recorded in golden matches request + cfg = golden["config"] + assert cfg["batch_size"] == batch + assert cfg["groups"] == g + assert cfg["use_bias"] == use_bias + # Output spatial from golden must match our shared calculate + assert via_golden.shape[2] == calculate_output_dim(h, k, s, p, 1) + assert via_golden.shape[3] == calculate_output_dim(w, k, s, p, 1) + + +@pytest.mark.parametrize( + "dummy", + [pytest.param(None, id="reference_sanity")], +) +def test_conv2d_reference_sanity(dummy): + """Sanity cross-checks and documentation of bf16 reference behavior (no HW). + + - Verifies generate_golden works for edge-ish sizes not in the main matrix. + - Documents that we rely on torch F.conv2d(bf16) as the reference (no + full ml_dtypes emulation like reduction sum/mean because conv MACs are + more complex). + - Quick reproducibility check: same seed -> identical golden across calls. + - Exercises conv2d_cpu directly with dilation=1 (the only supported value). + """ + torch.manual_seed(2026) + + # Reproducibility: two independent calls with same seed must match exactly + g1 = generate_golden_reference( + batch_size=2, + in_channels=8, + in_height=17, + in_width=19, + out_channels=4, + kernel_size=3, + stride=1, + padding=1, + groups=1, + use_bias=True, + seed=123, + ) + g2 = generate_golden_reference( + batch_size=2, + in_channels=8, + in_height=17, + in_width=19, + out_channels=4, + kernel_size=3, + stride=1, + padding=1, + groups=1, + use_bias=True, + seed=123, + ) + assert torch.equal(g1["input"], g2["input"]) + assert torch.equal(g1["weight"], g2["weight"]) + assert torch.equal(g1["bias"], g2["bias"]) + assert torch.equal(g1["output"], g2["output"]) + + # Direct conv2d_cpu sanity (covers a non-default spatial + stride + no bias) + x = g1["input"][:1] # take first batch element + w = g1["weight"] + direct_out = conv2d_cpu(x, w, bias=None, stride=2, padding=0, groups=1) + # Must have the shape predicted by the shared calculator + exp_h = calculate_output_dim(17, 3, 2, 0, 1) + exp_w = calculate_output_dim(19, 3, 2, 0, 1) + assert direct_out.shape == (1, 4, exp_h, exp_w) + + # bf16 vs "higher precision" reference drift note (for future tolerance tuning) + # We compute a quick fp32 reference for the same bf16-cast inputs to show + # the magnitude of bf16 rounding effect (not a test failure, just visibility). + x_fp32 = x.to(torch.float32) + w_fp32 = w.to(torch.float32) + fp32_ref = F.conv2d(x_fp32, w_fp32, bias=None, stride=2, padding=0, groups=1) + bf16_from_fp32 = fp32_ref.to(torch.bfloat16) + max_abs_drift = (bf16_from_fp32 - direct_out).abs().max().item() + # Drift is expected; we only log if "surprisingly large" for awareness. + if max_abs_drift > 0.5: + print( + f"[conv2d ref sanity] observed bf16-vs-fp32-ref drift={max_abs_drift:.4f} " + "(expected for bf16 conv; justifies 0.05 rel tol in HW tests)" + ) + # Always pass; this is informational only. + + +@pytest.mark.parametrize("dummy", [pytest.param(None, id="hw_tolerances")]) +def test_hw_tolerances_defaults(dummy): + """Default HW tolerances are centralized and stricter than HW_LOOSE.""" + from iron.operators.conv2d.tolerances import ( + HW_DEFAULT, + HW_LOOSE, + hw_tolerances, + ) + + t = hw_tolerances() + assert t is HW_DEFAULT + assert t.rel_tol < HW_LOOSE.rel_tol + assert t.abs_tol < HW_LOOSE.abs_tol + assert 0 < t.max_error_rate <= HW_LOOSE.max_error_rate + assert t.rel_tol < 1.0 and t.abs_tol > 0 + + +@pytest.mark.parametrize("dummy", [pytest.param(None, id="pack_weights_bias")]) +def test_pack_weights_with_bias_layout(dummy): + """Tile-interleaved W‖B pack matches design contract (groups==1, 2 cols).""" + import numpy as np + from ml_dtypes import bfloat16 + + from iron.operators.conv2d.design import pack_weights_with_bias + + oc, ic, kh, kw = 8, 4, 3, 3 + wpo = ic * kh * kw + w = np.arange(oc * wpo, dtype=np.float32).astype(bfloat16) + b = (np.arange(oc, dtype=np.float32) + 100).astype(bfloat16) + packed = pack_weights_with_bias( + w, + b, + out_channels=oc, + in_channels=ic, + groups=1, + kernel_h=kh, + kernel_w=kw, + num_columns=2, + is_depthwise=False, + tile_channels=4, + ) + # 2 cols × 1 tile × (4*wpo + 4 bias) + assert packed.shape[0] == oc * wpo + oc + # First tile: OCs 0..3 + assert np.array_equal(packed[: 4 * wpo], w[: 4 * wpo]) + assert np.array_equal(packed[4 * wpo : 4 * wpo + 4], b[:4]) + + +@pytest.mark.parametrize( + "dummy", [pytest.param(None, id="pack_weights_bias_grouped_2c")] +) +def test_pack_weights_with_bias_grouped_multicol(dummy): + """Grouped multi-col pack: OC blocks per column (groups % cols == 0).""" + import numpy as np + from ml_dtypes import bfloat16 + + from iron.operators.conv2d.design import ( + _resolve_num_columns, + pack_weights_with_bias, + ) + + # g=2, IC=8, OC=16 → 2 cols ⇒ g_per_col=1, ic_per_col=4, oc_per_col=8 + oc, ic, g, kh, kw = 16, 8, 2, 3, 3 + wpo = (ic // g) * kh * kw + assert _resolve_num_columns(2, oc, ic, g, False, max_cols=8) == 2 + assert _resolve_num_columns(3, oc, ic, g, False, max_cols=8) == 2 # clamp + assert _resolve_num_columns(8, oc, ic, g, False, max_cols=8) == 2 + w = np.arange(oc * wpo, dtype=np.float32).astype(bfloat16) + b = (np.arange(oc, dtype=np.float32) + 50).astype(bfloat16) + packed = pack_weights_with_bias( + w, + b, + out_channels=oc, + in_channels=ic, + groups=g, + kernel_h=kh, + kernel_w=kw, + num_columns=2, + is_depthwise=False, + tile_channels=8, # full oc_per_col + ) + assert packed.shape[0] == oc * wpo + oc + # Col0 tile: OC 0..7 weights then bias + assert np.array_equal(packed[: 8 * wpo], w[: 8 * wpo]) + assert np.array_equal(packed[8 * wpo : 8 * wpo + 8], b[:8]) + # Col1 tile: OC 8..15 + mid = 8 * wpo + 8 + assert np.array_equal(packed[mid : mid + 8 * wpo], w[8 * wpo :]) + assert np.array_equal(packed[mid + 8 * wpo :], b[8:]) + + +# Tests are pytest-only (AGENTS.md convention). diff --git a/iron/operators/conv2d/design.py b/iron/operators/conv2d/design.py new file mode 100644 index 00000000..9afe01c7 --- /dev/null +++ b/iron/operators/conv2d/design.py @@ -0,0 +1,1194 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +MLIR generation for AIE conv2d (AIE2 / AIE2P). + +Hard constraints (current design): +- Each compute tile has 2 input DMA channels: ObjectFifos are input + weight only. + When ``use_bias``, bias is **packed** after each weight tile + (``[W_tile ‖ B_tile]``) so apply_bias runs on-device without a 3rd DMA. + When ``use_bias`` is false, kernels run with apply_bias=0. +- L1 holds one in+weight(+bias)+out triple per iteration (budget + ``_L1_TRIPLE_BUDGET_BYTES``; FIFO depth 1 or 2 if 2× triple fits). +- groups==1: optional multi-col OC split + OC or H-strip tiling to fit L1. +- depthwise: channel split/tile across columns (no full-input broadcast). +- other groups>1: multi-col **group-block** split when ``groups % cols == 0`` + and the per-col IC/OC triple fits L1; else 1-col full triple or k>1 + host-pad H-strip. Multi-col grouped does not yet combine with H-strip. +- H-strip TAPs use leading size 1 so aiex does not treat strip count as + repeat_count; k>1 path may host-pad input and crop design OH/OW on host. +""" + +from ml_dtypes import bfloat16 +from pathlib import Path +import numpy as np +import argparse +import sys + +from aie.iron import Kernel, ObjectFifo, Program, Runtime, Worker +from aie.iron.device import NPU1, NPU2 +from aie.helpers.taplib.tap import TensorAccessPattern +from aie.iron.controlflow import range_ + +# Leave headroom under ~64KB L1 for stack/locks when depth=1 holds in+w+out. +_L1_TRIPLE_BUDGET_BYTES = 56 * 1024 +_BYTES_PER_BF16 = 2 + + +def _largest_divisor_fit(n: int, fits) -> int: + """Largest positive divisor of ``n`` for which ``fits(d)`` is true, else 1.""" + if n <= 0: + return 1 + if fits(n): + return n + for d in range(n - 1, 0, -1): + if n % d == 0 and fits(d): + return d + return 1 + + +def _choose_oc_tile( + out_channels: int, + input_elems: int, + weight_per_oc: int, + out_spatial: int, + l1_budget_bytes: int = _L1_TRIPLE_BUDGET_BYTES, +) -> int: + """Largest ``oc_tile`` dividing ``out_channels`` whose L1 triple fits. + + Triple = full input + weight tile + output tile (bf16). + Returns 1 if even a single OC does not fit; callers may then use H-strip + tiling or raise at construct time. + """ + + def fits(oc_t: int) -> bool: + elems = input_elems + oc_t * weight_per_oc + oc_t * out_spatial + return elems * _BYTES_PER_BF16 <= l1_budget_bytes + + return _largest_divisor_fit(out_channels, fits) + + +def _choose_channel_tile( + channels: int, + in_spatial: int, + out_spatial: int, + weight_per_c: int, + l1_budget_bytes: int = _L1_TRIPLE_BUDGET_BYTES, +) -> int: + """Largest channel tile for depthwise: tiles in+w+out together. + + Per-channel elems = in_spatial + weight_per_c + out_spatial (bf16). + """ + + def fits(c_t: int) -> bool: + elems = c_t * (in_spatial + weight_per_c + out_spatial) + return elems * _BYTES_PER_BF16 <= l1_budget_bytes + + return _largest_divisor_fit(channels, fits) + + +def _choose_h_tile_pointwise( + height: int, + in_channels: int, + width: int, + oc_per_col: int, + weight_per_oc: int, + l1_budget_bytes: int = _L1_TRIPLE_BUDGET_BYTES, +) -> int: + """Largest ``tile_h | height`` so the full ``oc_per_col`` triple fits. + + Prefers **num_oc_tiles=1** with H-strip spatial only. AIE DMA BDs require + positive strides, so we avoid multi-dim rebroadcast (stride 0) of input + across OC tiles or weights across spatial tiles. + + Pointwise: in = IC*th*W, weight = oc_per_col*weight_per_oc, out = oc*th*W. + Falls back to largest th where at least OC=1 fits (caller may still CE). + """ + + def fits_full_oc(th: int) -> bool: + elems = ( + in_channels * th * width + + oc_per_col * weight_per_oc + + oc_per_col * th * width + ) + return elems * _BYTES_PER_BF16 <= l1_budget_bytes + + th = _largest_divisor_fit(height, fits_full_oc) + if fits_full_oc(th): + return th + + def fits_min_oc(th: int) -> bool: + elems = in_channels * th * width + weight_per_oc + th * width + return elems * _BYTES_PER_BF16 <= l1_budget_bytes + + return _largest_divisor_fit(height, fits_min_oc) + + +def _rf_in_h(tile_oh: int, stride_h: int, kernel_h: int) -> int: + """Input rows needed for ``tile_oh`` output rows (pad=0, fixed RF).""" + return (max(1, tile_oh) - 1) * stride_h + kernel_h + + +def _extend_in_for_dma_even_out( + in_h: int, + in_w: int, + kernel_h: int, + kernel_w: int, + stride_h: int, + stride_w: int, + pad_h: int, + pad_w: int, +) -> tuple: + """Minimal bottom/right input growth so OH and OW are both even and >=1. + + Odd OH/OW blocks H-strip TAP dims (bf16 BD sizes must be even). Extra + input pixels are zeros on the host; valid crop is the un-extended out + spatial (op crops after NPU). Returns (in_h', in_w', out_h', out_w'). + """ + + def _out(h, w): + oh = (h + 2 * pad_h - kernel_h) // stride_h + 1 + ow = (w + 2 * pad_w - kernel_w) // stride_w + 1 + return oh, ow + + h, w = int(in_h), int(in_w) + for _ in range(h + w + 8): + oh, ow = _out(h, w) + if oh >= 1 and ow >= 1 and (oh % 2 == 0) and (ow % 2 == 0): + return h, w, oh, ow + if oh < 1 or (oh % 2 != 0): + h += 1 + elif ow < 1 or (ow % 2 != 0): + w += 1 + else: + h += 1 + oh, ow = _out(h, w) + return h, w, oh, ow + + +def _choose_h_tile_standard( + out_height: int, + in_channels: int, + padded_w: int, + oc_per_col: int, + weight_per_oc: int, + out_width: int, + kernel_h: int, + stride_h: int, + l1_budget_bytes: int = _L1_TRIPLE_BUDGET_BYTES, +) -> int: + """Largest ``tile_oh | out_height`` so full ``oc_per_col`` RF triple fits. + + Host-padded k>1 path: input strip height = + ``(tile_oh-1)*stride_h + kernel_h``, width = padded_w, pad=0 in kernel. + Prefers num_oc_tiles=1 (same DMA constraint as pointwise H-strip). + """ + + def fits_full_oc(toh: int) -> bool: + ih = _rf_in_h(toh, stride_h, kernel_h) + elems = ( + in_channels * ih * padded_w + + oc_per_col * weight_per_oc + + oc_per_col * toh * out_width + ) + return elems * _BYTES_PER_BF16 <= l1_budget_bytes + + th = _largest_divisor_fit(out_height, fits_full_oc) + if fits_full_oc(th): + return th + + def fits_min_oc(toh: int) -> bool: + ih = _rf_in_h(toh, stride_h, kernel_h) + elems = in_channels * ih * padded_w + weight_per_oc + toh * out_width + return elems * _BYTES_PER_BF16 <= l1_budget_bytes + + return _largest_divisor_fit(out_height, fits_min_oc) + + +def _plan_halo_h_strip( + in_height: int, + in_width: int, + true_out_height: int, + true_out_width: int, + in_channels: int, + oc_per_col: int, + weight_per_oc: int, + kernel_h: int, + kernel_w: int, + stride_h: int, + stride_w: int, + pad_h: int, + pad_w: int, + l1_budget_bytes: int = _L1_TRIPLE_BUDGET_BYTES, + max_extra: int = 16, +): + """Plan k>1 host-pad RF H-strip; may add bottom/right DMA pad. + + When natural padded spatial dims yield only odd DMA transfer sizes + (e.g. s2 p0 → OW=31, toh∈{1,31}), search a small bottom/right extra + zero-pad so design OH/OW admit even bf16 strip lengths. Host crops the + NPU output back to ``true_out_*``. + + Also searches **all** ``tile_oh | design_oh`` (large→small). Max L1-legal + toh can still violate ``aie.dma_bd`` size-dim u10 max (1023) when + ``in_h_tile * padded_w`` is large — e.g. toh=32 on 66-wide → 2244. + + Returns a dict on success:: + padded_h, padded_w, design_oh, design_ow, tile_oh, in_h_tile, + num_spatial, extra_h, extra_w + or ``None`` if no legal pure-H-strip plan (num_oc_tiles==1) fits L1 with + DMA-aligned strip sizes and BD-legal size dims. + """ + if oc_per_col <= 0 or true_out_height <= 0 or true_out_width <= 0: + return None + + # Prefer zero extra, then minimal total extra (eh,ew partitions of total). + candidates = [(0, 0)] + for total in range(1, max_extra + 1): + for eh in range(0, total + 1): + candidates.append((eh, total - eh)) + best = None + best_key = None + + for extra_h, extra_w in candidates: + padded_h = in_height + 2 * pad_h + extra_h + padded_w = in_width + 2 * pad_w + extra_w + if padded_h < kernel_h or padded_w < kernel_w: + continue + design_oh = (padded_h - kernel_h) // stride_h + 1 + design_ow = (padded_w - kernel_w) // stride_w + 1 + if design_oh < true_out_height or design_ow < true_out_width: + continue + if design_oh <= 0 or design_ow <= 0: + continue + + # large→small toh: first legal is max toh for this pad (break after). + for tile_oh in range(design_oh, 0, -1): + if design_oh % tile_oh != 0: + continue + num_spatial = design_oh // tile_oh + # Need multi-strip spatial tiling; even packet count for bf16 BDs. + if num_spatial <= 1 or (num_spatial % 2 != 0): + continue + in_h_tile = _rf_in_h(tile_oh, stride_h, kernel_h) + last_end = (num_spatial - 1) * tile_oh * stride_h + in_h_tile + if last_end > padded_h: + continue + out_strip = tile_oh * design_ow + in_strip = in_h_tile * padded_w + # bf16 BD granularity: even elem counts; u10 size dims ≤1023. + if (out_strip % 2 != 0) or (in_strip % 2 != 0): + continue + if ( + in_strip > 1023 + or out_strip > 1023 + or in_channels > 1023 + or oc_per_col > 1023 + or num_spatial > 1023 + ): + continue + in_tile = in_channels * in_h_tile * padded_w + out_tile = oc_per_col * tile_oh * design_ow + w_tile = oc_per_col * weight_per_oc + if not _l1_triple_fits(in_tile, w_tile, out_tile, l1_budget_bytes): + continue + + # Prefer: zero extra, then smaller total extra, larger tile_oh. + key = ( + extra_h + extra_w, + abs(extra_h - extra_w), + -tile_oh, + padded_h + padded_w, + ) + if best is None or key < best_key: + best_key = key + best = { + "padded_h": padded_h, + "padded_w": padded_w, + "design_oh": design_oh, + "design_ow": design_ow, + "tile_oh": tile_oh, + "in_h_tile": in_h_tile, + "num_spatial": num_spatial, + "extra_h": extra_h, + "extra_w": extra_w, + } + break # largest legal toh for this (extra_h, extra_w) + + return best + + +def _l1_triple_fits( + input_elems: int, + weight_elems: int, + output_elems: int, + l1_budget_bytes: int = _L1_TRIPLE_BUDGET_BYTES, +) -> bool: + """True if in+weight+out (bf16) fit the L1 triple budget.""" + return ( + input_elems + weight_elems + output_elems + ) * _BYTES_PER_BF16 <= l1_budget_bytes + + +def _bias_per_oc(use_bias: bool) -> int: + """Extra L1/L3 weight-buffer elems per OC (or depthwise channel) for packed bias.""" + return 1 if use_bias else 0 + + +def pack_weights_with_bias( + weight_flat, + bias_flat, + *, + out_channels: int, + in_channels: int, + groups: int, + kernel_h: int, + kernel_w: int, + num_columns: int, + is_depthwise: bool, + tile_channels: int, +): + """Host pack tile-interleaved ``[W_tile ‖ B_tile]`` for on-device bias. + + Matches design L3 weight buffer when ``use_bias``:: + for col: + for tile in column: + weights for tile_channels OCs (or depthwise channels) + bias for those channels + + ``tile_channels`` is ``oc_tile`` (groups==1 / grouped) or ``c_tile`` + (depthwise) — the same value design uses for L1 weight tiles. + """ + import numpy as np + + w = np.asarray(weight_flat).reshape(-1) + b = np.asarray(bias_flat).reshape(-1) + if b.shape[0] != out_channels: + raise ValueError(f"bias length {b.shape[0]} != out_channels {out_channels}") + weight_per_oc = (in_channels // groups) * kernel_h * kernel_w + expected_w = out_channels * weight_per_oc + if w.shape[0] != expected_w: + raise ValueError(f"weight length {w.shape[0]} != expected {expected_w}") + tc = max(1, int(tile_channels)) + cols = max(1, int(num_columns)) + parts = [] + + if is_depthwise: + if in_channels % cols != 0: + raise ValueError("depthwise pack requires in_channels % cols == 0") + c_per_col = in_channels // cols + if c_per_col % tc != 0: + raise ValueError( + f"depthwise tile_channels={tc} must divide c_per_col={c_per_col}" + ) + for i in range(cols): + c0 = i * c_per_col + for t in range(c_per_col // tc): + cs = c0 + t * tc + parts.append(w[cs * weight_per_oc : (cs + tc) * weight_per_oc]) + parts.append(b[cs : cs + tc]) + return np.concatenate(parts) + + # groups==1 multi-col OC split, or grouped multi-col group-block OC split. + # Grouped multi-col requires groups % cols == 0 so each column owns a + # contiguous block of groups (contiguous IC/OC in torch layout). + if out_channels % cols != 0: + raise ValueError("pack requires out_channels % cols == 0") + if groups > 1 and groups % cols != 0: + raise ValueError("grouped pack requires groups % cols == 0") + oc_per_col = out_channels // cols + if oc_per_col % tc != 0: + raise ValueError(f"tile_channels={tc} must divide oc_per_col={oc_per_col}") + for i in range(cols): + o0 = i * oc_per_col + for t in range(oc_per_col // tc): + os_ = o0 + t * tc + parts.append(w[os_ * weight_per_oc : (os_ + tc) * weight_per_oc]) + parts.append(b[os_ : os_ + tc]) + return np.concatenate(parts) + + +def _resolve_num_columns( + requested: int, + out_channels: int, + in_channels: int, + groups: int, + is_depthwise: bool, + max_cols: int, +) -> int: + """Clamp column count for legal OC/channel/group splits and device limits.""" + n = max(1, int(requested) if requested is not None else 1) + n = min(n, max_cols) + if is_depthwise: + while n > 1 and in_channels % n != 0: + n -= 1 + return n + if groups == 1: + while n > 1 and out_channels % n != 0: + n -= 1 + return n + # Non-depthwise grouped: multi-col when groups (hence IC/OC) divide n. + while n > 1 and (groups % n != 0 or in_channels % n != 0 or out_channels % n != 0): + n -= 1 + return n + + +def my_conv2d( + dev, + N, # batch size + in_channels, + in_height, + in_width, + out_channels, + out_height, + out_width, + kernel_h, + kernel_w, + stride_h, + stride_w, + pad_h, + pad_w, + groups, + use_bias, + num_columns, + tile_size, + trace_size, +): + """ + Generate MLIR for 2D convolution (L1 tiles + multi-col OC/channel split). + + When ``use_bias``, weights and bias are packed into one L3 buffer as + tile-interleaved ``[W_tile ‖ B_tile]`` (still one weight ObjectFifo — no + third DMA). Columns: groups==1 OC-split, depthwise channel-split, and + non-depthwise grouped **group-block** split when ``groups % cols == 0`` + and the per-col triple fits L1; otherwise clamped toward 1. + """ + dtype = bfloat16 + + _ = (tile_size, trace_size) + bias_extra = _bias_per_oc(bool(use_bias)) + + # Device column cap from target model (NPU1.cols==4, NPU2.cols==8). + max_cols = getattr(dev, "cols", None) or 4 + + input_size = N * in_channels * in_height * in_width + weight_only_size = out_channels * in_channels // groups * kernel_h * kernel_w + # L3 weight BO: pure weights, or packed W‖B (one bias per OC/channel). + weight_size = weight_only_size + (out_channels if bias_extra else 0) + output_size = N * out_channels * out_height * out_width + in_spatial = in_height * in_width + out_spatial = out_height * out_width + weight_per_oc = (in_channels // groups) * kernel_h * kernel_w + # Storage elems per OC in the packed weight stream (weights + optional bias). + weight_store_per_oc = weight_per_oc + bias_extra + + input_ty = np.ndarray[(input_size,), np.dtype[dtype]] + weight_ty = np.ndarray[(weight_size,), np.dtype[dtype]] + output_ty = np.ndarray[(output_size,), np.dtype[dtype]] + + # Variant selection (must match C++ symbols in aie_kernels/*/conv2d.cc). + is_depthwise = groups == in_channels and groups == out_channels + is_pointwise = (not is_depthwise) and kernel_h == 1 and kernel_w == 1 + if is_depthwise: + kernel_name = "depthwise_conv2d_bf16_vector" + elif is_pointwise: + kernel_name = "pointwise_conv2d_bf16_vector" + else: + kernel_name = "conv2d_bf16_vector" + + num_columns = _resolve_num_columns( + num_columns, out_channels, in_channels, groups, is_depthwise, max_cols + ) + + # --- Per-column tile selection + multi-col split sizes -------------------- + # rebroadcast_input: full input OF packet, repeated per tile (groups==1). + # depthwise_split: per-col channel blocks for in/w/out (no full-input broadcast). + # grouped_split: per-col IC/OC group blocks (non-DW groups>1 multi-col). + # spatial_h_tiling: H-strip when full input exceeds L1 (pointwise or k>1). + # spatial_halo_pad: k>1 host-padded RF strips (kernel pad=0; L3 input padded). + rebroadcast_input = False + depthwise_split = False + grouped_split = False + spatial_h_tiling = False + spatial_halo_pad = False + tile_h = in_height # output strip height when spatial; else full in/out H + in_h_tile = in_height # input strip height (RF size when spatial_halo_pad) + padded_h = in_height + padded_w = in_width + num_spatial = 1 + num_oc_tiles = 1 + # Per-column tensor footprints for TAPs (bytes/elems along OC or channel axis). + weight_elems_per_col = weight_size + output_elems_per_col = output_size + input_elems_per_col = input_size + # Kernel scalar dims (overridden for grouped multi-col mini-convs). + kernel_in_channels = in_channels + kernel_groups = groups + + # weight_only_tile_elems: pure weights in each L1 packet (kernel weight ptr). + # weight_tile_elems: packet size including trailing packed bias when enabled. + weight_only_tile_elems = 0 + + if is_depthwise: + # Split channels across columns; tile within each column. + c_per_col = in_channels // num_columns + c_tile = _choose_channel_tile( + c_per_col, in_spatial, out_spatial, weight_store_per_oc + ) + if c_per_col % c_tile != 0: + c_tile = c_per_col + num_tiles = c_per_col // c_tile + num_oc_tiles = num_tiles + input_tile_elems = N * c_tile * in_spatial + weight_only_tile_elems = c_tile * weight_per_oc + weight_tile_elems = c_tile * weight_store_per_oc + output_tile_elems = N * c_tile * out_spatial + kernel_channels = c_tile + oc_tile = c_tile + depthwise_split = True + input_elems_per_col = N * c_per_col * in_spatial + weight_elems_per_col = c_per_col * weight_store_per_oc + output_elems_per_col = N * c_per_col * out_spatial + elif groups == 1: + # OC split across columns; OC tile within each column. + # If full input still exceeds L1 → pointwise H-strip or k>1 host-pad RF. + oc_per_col = out_channels // num_columns + oc_tile = _choose_oc_tile( + oc_per_col, input_size, weight_store_per_oc, out_spatial + ) + if oc_per_col % oc_tile != 0: + oc_tile = oc_per_col + full_fits = _l1_triple_fits( + input_size, oc_tile * weight_store_per_oc, N * oc_tile * out_spatial + ) + if (not full_fits) and is_pointwise: + # Pointwise H-strip: prefer full oc_per_col in L1 (num_oc=1) + # so TAPs need no stride-0 rebroadcast (illegal on aie.dma_bd). + tile_h = _choose_h_tile_pointwise( + in_height, in_channels, in_width, oc_per_col, weight_store_per_oc + ) + if in_height % tile_h != 0: + tile_h = in_height + num_spatial = max(1, in_height // tile_h) + in_h_tile = tile_h + in_tile_elems_base = N * in_channels * tile_h * in_width + out_tile_sp = tile_h * out_width + # Prefer full OC block when it fits with this tile_h. + if _l1_triple_fits( + in_tile_elems_base, + oc_per_col * weight_store_per_oc, + N * oc_per_col * out_tile_sp, + ): + oc_tile = oc_per_col + else: + oc_tile = _choose_oc_tile( + oc_per_col, in_tile_elems_base, weight_store_per_oc, out_tile_sp + ) + if oc_per_col % oc_tile != 0: + oc_tile = oc_per_col + num_oc_tiles = oc_per_col // oc_tile if oc_tile else 1 + # Multi-dim spatial TAPs only for pure H-strip (no OC rebroadcast). + if num_oc_tiles != 1: + # Cannot legally TAP-rebroadcast; keep full-input path (will + # op._validate_l1_fit raises when the min tile cannot fit. + tile_h = in_height + in_h_tile = in_height + num_spatial = 1 + oc_tile = _choose_oc_tile( + oc_per_col, input_size, weight_store_per_oc, out_spatial + ) + if oc_per_col % oc_tile != 0: + oc_tile = oc_per_col + input_tile_elems = input_size + weight_only_tile_elems = oc_tile * weight_per_oc + weight_tile_elems = oc_tile * weight_store_per_oc + output_tile_elems = N * oc_tile * out_spatial + spatial_h_tiling = False + else: + spatial_h_tiling = num_spatial > 1 + input_tile_elems = in_tile_elems_base + weight_only_tile_elems = oc_tile * weight_per_oc + weight_tile_elems = oc_tile * weight_store_per_oc + output_tile_elems = N * oc_tile * out_tile_sp + elif not full_fits: + # Standard k>1 H-strip: host zero-pads (conv pad + optional + # bottom/right DMA pad); kernel pad=0 with fixed RF strip height; + # overlapping input TAPs. May use design_oh/ow > true out (crop). + plan = _plan_halo_h_strip( + in_height, + in_width, + out_height, + out_width, + in_channels, + oc_per_col, + weight_store_per_oc, + kernel_h, + kernel_w, + stride_h, + stride_w, + pad_h, + pad_w, + ) + if plan is not None: + spatial_h_tiling = True + spatial_halo_pad = True + padded_h = plan["padded_h"] + padded_w = plan["padded_w"] + design_oh = plan["design_oh"] + design_ow = plan["design_ow"] + tile_oh = plan["tile_oh"] + in_h_tile = plan["in_h_tile"] + num_spatial = plan["num_spatial"] + tile_h = tile_oh + oc_tile = oc_per_col + num_oc_tiles = 1 + in_tile_elems_base = N * in_channels * in_h_tile * padded_w + out_tile_sp = tile_oh * design_ow + input_tile_elems = in_tile_elems_base + weight_only_tile_elems = oc_tile * weight_per_oc + weight_tile_elems = oc_tile * weight_store_per_oc + output_tile_elems = N * oc_tile * out_tile_sp + # L3: padded input; output may be design spatial (host crops). + input_size = N * in_channels * padded_h * padded_w + input_ty = np.ndarray[(input_size,), np.dtype[dtype]] + if design_oh != out_height or design_ow != out_width: + out_height = design_oh + out_width = design_ow + out_spatial = out_height * out_width + output_size = N * out_channels * out_spatial + output_ty = np.ndarray[(output_size,), np.dtype[dtype]] + else: + tile_h = out_height + in_h_tile = in_height + num_spatial = 1 + padded_h = in_height + padded_w = in_width + oc_tile = _choose_oc_tile( + oc_per_col, input_size, weight_store_per_oc, out_spatial + ) + if oc_per_col % oc_tile != 0: + oc_tile = oc_per_col + input_tile_elems = input_size + weight_only_tile_elems = oc_tile * weight_per_oc + weight_tile_elems = oc_tile * weight_store_per_oc + output_tile_elems = N * oc_tile * out_spatial + spatial_h_tiling = False + spatial_halo_pad = False + else: + input_tile_elems = input_size + weight_only_tile_elems = oc_tile * weight_per_oc + weight_tile_elems = oc_tile * weight_store_per_oc + output_tile_elems = N * oc_tile * out_spatial + + num_oc_tiles = oc_per_col // oc_tile if oc_tile else 1 + num_tiles = num_spatial * num_oc_tiles + if not spatial_h_tiling: + rebroadcast_input = num_oc_tiles > 1 + kernel_channels = in_channels + weight_elems_per_col = oc_per_col * weight_store_per_oc + output_elems_per_col = N * oc_per_col * out_spatial + if weight_only_tile_elems == 0: + weight_only_tile_elems = oc_tile * weight_per_oc + weight_tile_elems = oc_tile * weight_store_per_oc + else: + # Non-depthwise grouped: multi-col group-block split when L1 allows, + # else 1-col full tensor or k>1 host-pad H-strip. + # Each column owns groups/cols contiguous groups → contiguous IC/OC. + ic_per_col = in_channels // num_columns + oc_per_col = out_channels // num_columns + groups_per_col = groups // num_columns + input_elems_per_col = N * ic_per_col * in_spatial + weight_elems_per_col = oc_per_col * weight_store_per_oc + output_elems_per_col = N * oc_per_col * out_spatial + per_col_fits = _l1_triple_fits( + input_elems_per_col, weight_elems_per_col, output_elems_per_col + ) + if per_col_fits: + # Full per-col triple (num_tiles=1). Multi-col uses grouped_split TAPs. + oc_tile = oc_per_col + num_oc_tiles = 1 + num_tiles = 1 + input_tile_elems = input_elems_per_col + weight_only_tile_elems = oc_per_col * weight_per_oc + weight_tile_elems = weight_elems_per_col + output_tile_elems = output_elems_per_col + kernel_channels = ic_per_col + kernel_in_channels = ic_per_col + kernel_groups = groups_per_col + grouped_split = num_columns > 1 + if not grouped_split: + # 1-col: keep global sizes (same footprint; simpler TAPs). + input_tile_elems = input_size + weight_only_tile_elems = weight_only_size + weight_tile_elems = weight_size + output_tile_elems = output_size + kernel_in_channels = in_channels + kernel_groups = groups + kernel_channels = in_channels + weight_elems_per_col = weight_size + output_elems_per_col = output_size + input_elems_per_col = input_size + else: + # Multi-col H-strip for groups is not implemented: clamp to 1-col. + num_columns = 1 + ic_per_col = in_channels + oc_per_col = out_channels + groups_per_col = groups + kernel_in_channels = in_channels + kernel_groups = groups + kernel_channels = in_channels + weight_elems_per_col = weight_size + output_elems_per_col = output_size + input_elems_per_col = input_size + if _l1_triple_fits(input_size, weight_size, output_size): + oc_tile = out_channels + num_tiles = 1 + input_tile_elems = input_size + weight_only_tile_elems = weight_only_size + weight_tile_elems = weight_size + output_tile_elems = output_size + else: + plan = _plan_halo_h_strip( + in_height, + in_width, + out_height, + out_width, + in_channels, + oc_per_col, + weight_store_per_oc, + kernel_h, + kernel_w, + stride_h, + stride_w, + pad_h, + pad_w, + ) + if plan is not None: + spatial_h_tiling = True + spatial_halo_pad = True + padded_h = plan["padded_h"] + padded_w = plan["padded_w"] + design_oh = plan["design_oh"] + design_ow = plan["design_ow"] + tile_oh = plan["tile_oh"] + in_h_tile = plan["in_h_tile"] + num_spatial = plan["num_spatial"] + tile_h = tile_oh + oc_tile = oc_per_col + num_oc_tiles = 1 + num_tiles = num_spatial + in_tile_elems_base = N * in_channels * in_h_tile * padded_w + out_tile_sp = tile_oh * design_ow + input_tile_elems = in_tile_elems_base + weight_only_tile_elems = oc_tile * weight_per_oc + weight_tile_elems = oc_tile * weight_store_per_oc + output_tile_elems = N * oc_tile * out_tile_sp + input_size = N * in_channels * padded_h * padded_w + input_ty = np.ndarray[(input_size,), np.dtype[dtype]] + if design_oh != out_height or design_ow != out_width: + out_height = design_oh + out_width = design_ow + out_spatial = out_height * out_width + output_size = N * out_channels * out_spatial + output_ty = np.ndarray[(output_size,), np.dtype[dtype]] + weight_elems_per_col = weight_tile_elems + output_elems_per_col = output_size + else: + oc_tile = out_channels + num_tiles = 1 + input_tile_elems = input_size + weight_only_tile_elems = weight_only_size + weight_tile_elems = weight_size + output_tile_elems = output_size + + # FIFO element types = per-iteration L1 footprints. + input_tile_ty = np.ndarray[ + (input_tile_elems if input_tile_elems > 0 else 1,), np.dtype[dtype] + ] + weight_tile_ty = np.ndarray[ + (weight_tile_elems if weight_tile_elems > 0 else 1,), np.dtype[dtype] + ] + output_tile_ty = np.ndarray[ + (output_tile_elems if output_tile_elems > 0 else 1,), np.dtype[dtype] + ] + + # depth=2 when 2x triple fits; else depth=1 (ping-pong would blow L1). + triple_bytes = ( + input_tile_elems + weight_tile_elems + output_tile_elems + ) * _BYTES_PER_BF16 + fifodepth = 1 if triple_bytes * 2 > _L1_TRIPLE_BUDGET_BYTES else 2 + + of_ins = [ + ObjectFifo(input_tile_ty, name=f"in_{i}", depth=fifodepth) + for i in range(num_columns) + ] + of_weights = [ + ObjectFifo(weight_tile_ty, name=f"w_{i}", depth=fifodepth) + for i in range(num_columns) + ] + of_outs = [ + ObjectFifo(output_tile_ty, name=f"out_{i}", depth=fifodepth) + for i in range(num_columns) + ] + + # On-device packed bias (weights‖bias in one ObjectFifo); no third DMA. + apply_bias = 1 if bias_extra else 0 + + if kernel_name == "depthwise_conv2d_bf16_vector": + # Mini depthwise over c_tile channels (or full when num_tiles==1). + kernel_int_types = [np.int32] * 13 + kernel_call_scalars = [ + N, + kernel_channels, + in_height, + in_width, + out_height, + out_width, + kernel_h, + kernel_w, + stride_h, + stride_w, + pad_h, + pad_w, + apply_bias, + ] + elif kernel_name == "pointwise_conv2d_bf16_vector": + # Mini pointwise over oc_tile out-channels; height may be H-strip. + kernel_int_types = [np.int32] * 6 + kernel_call_scalars = [ + N, + kernel_in_channels, + oc_tile, + tile_h, + in_width, + apply_bias, + ] + else: + # Standard mini-conv: out_channels = oc_tile when groups==1 tiled. + # Grouped multi-col: kernel sees local IC/OC/groups per column. + # Halo H-strip: strip-local spatial dims + pad=0 (host supplies pad). + k_in_h = in_h_tile if spatial_halo_pad else in_height + k_in_w = padded_w if spatial_halo_pad else in_width + k_out_h = tile_h if spatial_halo_pad else out_height + k_pad_h = 0 if spatial_halo_pad else pad_h + k_pad_w = 0 if spatial_halo_pad else pad_w + kernel_int_types = [np.int32] * 15 + kernel_call_scalars = [ + N, + kernel_in_channels, + k_in_h, + k_in_w, + oc_tile, + k_out_h, + out_width, + kernel_h, + kernel_w, + stride_h, + stride_w, + k_pad_h, + k_pad_w, + kernel_groups, + apply_bias, + ] + + # 4th buffer arg kept for ABI. Kernels with apply_bias=1 read bias from the + # tail of the weight tile (packed W‖B); dummy pointer is unused then. + bias_arg_ty = input_tile_ty + + conv2d_kernel = Kernel( + kernel_name, + "conv2d.o", + [input_tile_ty, weight_tile_ty, output_tile_ty, bias_arg_ty] + kernel_int_types, + ) + + def core_body(of_in, of_w, of_out, conv_kernel): + # One mini-conv per tile (num_tiles==1 => single full-tensor iter). + # Spatial H-strip: num_tiles == num_spatial (num_oc_tiles==1); weights + # rebroadcast via outermost TAP dim stride 0. + for _ in range_(num_tiles): + elem_in = of_in.acquire(1) + elem_w = of_w.acquire(1) + elem_out = of_out.acquire(1) + # Dummy bias pointer (packed bias uses weight tail when apply_bias=1). + elem_bias = elem_in + conv_kernel(elem_in, elem_w, elem_out, elem_bias, *kernel_call_scalars) + of_in.release(1) + of_w.release(1) + of_out.release(1) + + my_workers = [ + Worker( + core_body, + [ + of_ins[i].cons(), + of_weights[i].cons(), + of_outs[i].prod(), + conv2d_kernel, + ], + ) + for i in range(num_columns) + ] + + # --- TAPs: per-column offsets; multi-packet within column when tiling ------ + if spatial_h_tiling: + # H-strip (pointwise or k>1 host-pad RF), num_oc_tiles==1. + # CRITICAL (aiex.shim_dma_single_bd_task): sizes[0] becomes + # repeat_count=sizes[0]-1 and transfer_len=prod(sizes[-3:]). + # For strided multi-packet, put a leading 1 so repeat_count=0 and + # transfer_len covers all strips (one BD, no BD-ID blowup). + # Weight rebroadcast uses leading num_spatial + stride 0. + if spatial_halo_pad: + # Overlapping RF strips on host-padded NCHW: step tile_oh * sh rows. + strip_elems = in_h_tile * padded_w + strip_step = tile_h * stride_h * padded_w + ch_plane = padded_h * padded_w + else: + # Pointwise: non-overlapping equal in/out H strips. + strip_elems = tile_h * in_width + strip_step = strip_elems + ch_plane = in_height * in_width + out_strip = tile_h * out_width + input_taps = [ + TensorAccessPattern( + (1, input_size), + 0, + [1, num_spatial, in_channels, strip_elems], + [0, strip_step, ch_plane, 1], + ) + for _ in range(num_columns) + ] + weight_taps = [ + TensorAccessPattern( + (1, weight_size), + i * weight_elems_per_col, + [num_spatial, 1, 1, weight_tile_elems], + [0, 0, 0, 1], + ) + for i in range(num_columns) + ] + output_taps = [ + TensorAccessPattern( + (1, output_size), + i * output_elems_per_col, + [1, num_spatial, oc_tile, out_strip], + [0, out_strip, out_height * out_width, 1], + ) + for i in range(num_columns) + ] + elif depthwise_split or grouped_split: + # Channel/group blocks: in/w/out offset by column * elems_per_col. + # Depthwise: per-col channels. Grouped multi-col: per-col IC/OC groups. + input_taps = [ + TensorAccessPattern( + (1, input_size), + i * input_elems_per_col, + [1, 1, 1, input_elems_per_col], + [0, 0, 0, 1], + ) + for i in range(num_columns) + ] + weight_taps = [ + TensorAccessPattern( + (1, weight_size), + i * weight_elems_per_col, + [1, 1, 1, weight_elems_per_col], + [0, 0, 0, 1], + ) + for i in range(num_columns) + ] + output_taps = [ + TensorAccessPattern( + (1, output_size), + i * output_elems_per_col, + [1, 1, 1, output_elems_per_col], + [0, 0, 0, 1], + ) + for i in range(num_columns) + ] + elif rebroadcast_input: + # Full input rebroadcast once per OC tile (same on every column). + input_taps = [ + TensorAccessPattern( + (1, input_size), + 0, + [num_tiles, 1, 1, input_size], + [0, 0, 0, 1], + ) + for _ in range(num_columns) + ] + weight_taps = [ + TensorAccessPattern( + (1, weight_size), + i * weight_elems_per_col, + [1, 1, 1, weight_elems_per_col], + [0, 0, 0, 1], + ) + for i in range(num_columns) + ] + output_taps = [ + TensorAccessPattern( + (1, output_size), + i * output_elems_per_col, + [1, 1, 1, output_elems_per_col], + [0, 0, 0, 1], + ) + for i in range(num_columns) + ] + else: + # Single full-input transfer per column (num_tiles==1 groups==1 or grouped). + input_taps = [ + TensorAccessPattern( + (1, input_size), + 0, + [1, 1, 1, input_size], + [0, 0, 0, 1], + ) + for _ in range(num_columns) + ] + weight_taps = [ + TensorAccessPattern( + (1, weight_size), + i * weight_elems_per_col, + [1, 1, 1, weight_elems_per_col], + [0, 0, 0, 1], + ) + for i in range(num_columns) + ] + output_taps = [ + TensorAccessPattern( + (1, output_size), + i * output_elems_per_col, + [1, 1, 1, output_elems_per_col], + [0, 0, 0, 1], + ) + for i in range(num_columns) + ] + + rt = Runtime() + # Always 3 host buffers: in, weight(+packed bias), out. ≤2 input DMAs. + with rt.sequence(input_ty, weight_ty, output_ty) as (A, W, C): + rt.start(*my_workers) + tg = rt.task_group() + for i in range(num_columns): + rt.fill(of_ins[i].prod(), A, input_taps[i], task_group=tg) + for i in range(num_columns): + rt.fill(of_weights[i].prod(), W, weight_taps[i], task_group=tg) + for i in range(num_columns): + rt.drain( + of_outs[i].cons(), + C, + output_taps[i], + wait=True, + task_group=tg, + ) + rt.finish_task_group(tg) + + return Program(dev, rt).resolve_program() + + +if __name__ == "__main__": + + def str_to_device(device: str): + if device == "npu": + return NPU1() + elif device == "npu2": + return NPU2() + else: + raise ValueError(f"Device name {device} is unknown.") + + p = argparse.ArgumentParser() + p.add_argument( + "-d", + "--dev", + required=True, + dest="device", + help="AIE Device (npu or npu2)", + type=str_to_device, + ) + p.add_argument("-N", "--batch", type=int, default=1, help="Batch size") + p.add_argument( + "-ic", "--in-channels", type=int, required=True, help="Input channels" + ) + p.add_argument("-ih", "--in-height", type=int, required=True, help="Input height") + p.add_argument("-iw", "--in-width", type=int, required=True, help="Input width") + p.add_argument( + "-oc", "--out-channels", type=int, required=True, help="Output channels" + ) + p.add_argument("-kh", "--kernel-h", type=int, default=3, help="Kernel height") + p.add_argument("-kw", "--kernel-w", type=int, default=3, help="Kernel width") + p.add_argument("-sh", "--stride-h", type=int, default=1, help="Stride height") + p.add_argument("-sw", "--stride-w", type=int, default=1, help="Stride width") + p.add_argument("-ph", "--pad-h", type=int, default=0, help="Padding height") + p.add_argument("-pw", "--pad-w", type=int, default=0, help="Padding width") + p.add_argument("-g", "--groups", type=int, default=1, help="Number of groups") + p.add_argument("--use-bias", action="store_true", help="Use bias (host-side)") + p.add_argument( + "-co", + "--columns", + type=int, + default=1, + help="AIE columns (OC/channel split; clamped if not divisible)", + ) + p.add_argument("-ts", "--tile-size", type=int, default=1024, help="Tile size") + p.add_argument("-t", "--trace-size", type=int, default=0, help="Trace size") + p.add_argument( + "--output-file-path", + "-o", + type=str, + help="Output file path for the generated MLIR module", + ) + + opts = p.parse_args(sys.argv[1:]) + + dev = opts.device + N = opts.batch + in_channels = opts.in_channels + in_height = opts.in_height + in_width = opts.in_width + out_channels = opts.out_channels + kernel_h = opts.kernel_h + kernel_w = opts.kernel_w + stride_h = opts.stride_h + stride_w = opts.stride_w + pad_h = opts.pad_h + pad_w = opts.pad_w + groups = opts.groups + use_bias = opts.use_bias + columns = opts.columns + tile_size = opts.tile_size + trace_size = opts.trace_size + + if isinstance(dev, NPU1) and columns > 4: + raise ValueError("[ERROR] NPU device cannot allocate more than 4 columns") + elif isinstance(dev, NPU2) and columns > 8: + raise ValueError("[ERROR] NPU2 device cannot allocate more than 8 columns") + + out_height = (in_height + 2 * pad_h - kernel_h) // stride_h + 1 + out_width = (in_width + 2 * pad_w - kernel_w) // stride_w + 1 + + module = my_conv2d( + dev, + N, + in_channels, + in_height, + in_width, + out_channels, + out_height, + out_width, + kernel_h, + kernel_w, + stride_h, + stride_w, + pad_h, + pad_w, + groups, + use_bias, + columns, + tile_size, + trace_size, + ) + + output_file_path = Path(opts.output_file_path) + with open(output_file_path, "w") as f: + f.write(str(module)) diff --git a/iron/operators/conv2d/op.py b/iron/operators/conv2d/op.py new file mode 100644 index 00000000..916df4fb --- /dev/null +++ b/iron/operators/conv2d/op.py @@ -0,0 +1,943 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +AIE 2D Convolution Operator (AIE2 / AIE2P, bfloat16). + +Configurable kernel_size, stride, padding, groups (incl. depthwise). +Dilation is fixed to 1. Bias is packed on-device as ``weights‖bias`` in the +weight ObjectFifo (still ≤2 input DMAs: input + packed weight); kernels use +``apply_bias=1``. Construct-time checks enforce column policy and L1 triple +budget (including packed bias) via AIEOperatorConstraintError. +""" + +import torch +import numpy as np +from ml_dtypes import bfloat16 +from pathlib import Path +from typing import Tuple, Union, Optional, Callable, Any + +import aie.utils as aie_utils +from aie.utils.npukernel import NPUKernel +from aie.utils.hostruntime.xrtruntime.tensor import XRTTensor + +from iron.common import ( + AIEOperatorBase, + AIEOperatorConstraintError, + XclbinArtifact, + InstsBinArtifact, + KernelObjectArtifact, + SourceArtifact, + PythonGeneratedMLIRArtifact, + AIERuntimeArgSpec, + DesignGenerator, +) + +# Shared L1 / column policy with design.py (single source of truth). +from iron.operators.conv2d.design import ( + _BYTES_PER_BF16, + _L1_TRIPLE_BUDGET_BYTES, + _bias_per_oc, + _choose_channel_tile, + _choose_h_tile_pointwise, + _choose_h_tile_standard, + _choose_oc_tile, + _l1_triple_fits, + _plan_halo_h_strip, + _resolve_num_columns, + _rf_in_h, + pack_weights_with_bias, +) + + +class AIEConv2d(AIEOperatorBase): + """AIE-accelerated 2D convolution (bf16, AIE2 / AIE2P). + + Supports general k/stride/pad, groups (incl. depthwise), and on-device packed + bias (``[W_tile‖B_tile]``, ≤2 input DMAs). Dilation is only ``(1, 1)``. + Construct-time ``AIEOperatorConstraintError`` when dims, groups, or L1 plan + are illegal. Multi-col: OC-split, depthwise channel-split, or group-block + when ``groups % cols == 0``. Batch N is looped on the host over N=1 MLIR. + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: Union[int, Tuple[int, int]], + stride: Union[int, Tuple[int, int]] = 1, + padding: Union[int, Tuple[int, int]] = 0, + dilation: Union[int, Tuple[int, int]] = 1, + groups: int = 1, + use_bias: bool = True, + in_height: int = 32, + in_width: int = 32, + num_aie_columns: int = None, + tile_size: int = None, + context=None, + ): + """Build a specialized conv2d for the given channels and spatial size. + + ``in_height``/``in_width`` are compile-time; bias packs into the weight + DMA buffer when ``use_bias`` is true. + """ + self.in_channels = in_channels + self.out_channels = out_channels + + if isinstance(kernel_size, int): + kernel_size = (kernel_size, kernel_size) + if isinstance(stride, int): + stride = (stride, stride) + if isinstance(padding, int): + padding = (padding, padding) + if isinstance(dilation, int): + dilation = (dilation, dilation) + + self.kernel_size = kernel_size + self.stride = stride + self.padding = padding + self.dilation = dilation + self.groups = groups + self.use_bias = use_bias + self.in_height = in_height + self.in_width = in_width + + if in_channels <= 0 or out_channels <= 0: + raise AIEOperatorConstraintError( + f"AIEConv2d requires positive in_channels/out_channels, " + f"got in_channels={in_channels}, out_channels={out_channels}" + ) + if in_height <= 0 or in_width <= 0: + raise AIEOperatorConstraintError( + f"AIEConv2d requires positive in_height/in_width, " + f"got {in_height}x{in_width}" + ) + if dilation != (1, 1): + raise AIEOperatorConstraintError( + f"AIEConv2d only supports dilation=(1, 1), got {dilation}" + ) + if groups <= 0: + raise AIEOperatorConstraintError( + f"AIEConv2d requires groups >= 1, got {groups}" + ) + if in_channels % groups != 0: + raise AIEOperatorConstraintError( + f"AIEConv2d in_channels ({in_channels}) must be divisible by " + f"groups ({groups})" + ) + if out_channels % groups != 0: + raise AIEOperatorConstraintError( + f"AIEConv2d out_channels ({out_channels}) must be divisible by " + f"groups ({groups})" + ) + + self.out_height = ( + in_height + 2 * self.padding[0] - self.kernel_size[0] + ) // self.stride[0] + 1 + self.out_width = ( + in_width + 2 * self.padding[1] - self.kernel_size[1] + ) // self.stride[1] + 1 + if self.out_height <= 0 or self.out_width <= 0: + raise AIEOperatorConstraintError( + f"AIEConv2d produced non-positive output spatial size " + f"{self.out_height}x{self.out_width} from " + f"in={in_height}x{in_width}, kernel={self.kernel_size}, " + f"stride={self.stride}, padding={self.padding}" + ) + + if tile_size is None: + tile_size = 2048 + if num_aie_columns is None: + num_aie_columns = 1 + if int(num_aie_columns) < 1: + raise AIEOperatorConstraintError( + f"AIEConv2d num_aie_columns must be >= 1, got {num_aie_columns}" + ) + + self.tile_size = tile_size + self.num_aie_columns = int(num_aie_columns) + self.requested_num_columns = self.num_aie_columns + # Match design.py _resolve_num_columns. Device max_cols is applied in + # set_up_artifacts (and re-validated for L1 after the final clamp). + is_depthwise = groups == in_channels and groups == out_channels + self.is_depthwise = is_depthwise + # Construct-time: allow up to NPU2 max; set_up_artifacts tightens further. + # Grouped multi-col may further drop columns when per-col L1 does not fit + # (design then uses 1-col full/H-strip); keep host pack + design in sync. + self.effective_num_columns = self._resolve_columns_for_l1( + self.num_aie_columns, max_cols=8 + ) + self._validate_l1_fit(self.effective_num_columns) + + self.bias_size = out_channels if use_bias else 0 + + # Flattened N=1 sizes (batch looped in forward); used by get_arg_spec / forward. + self.input_size = in_channels * in_height * in_width + self.weight_size = ( + out_channels + * (in_channels // groups) + * self.kernel_size[0] + * self.kernel_size[1] + ) + self.output_size = out_channels * self.out_height * self.out_width + + self.xclbin_artifact = None + self.insts_artifact = None + self.weight_buffer = None + self.bias_buffer = None + # Cached NPU callable (invalidated on compile). + self._callable: Callable[..., Any] | None = None + + AIEOperatorBase.__init__(self, context=context) + + def _is_pointwise(self) -> bool: + return ( + (not self.is_depthwise) + and self.kernel_size[0] == 1 + and self.kernel_size[1] == 1 + ) + + def _resolve_columns_for_l1(self, requested: int, max_cols: int) -> int: + """Divisibility clamp, then drop columns until L1 policy accepts. + + For non-depthwise grouped multi-col, design requires the **per-col** + IC/OC triple to fit (no multi-col H-strip yet). If it does not, fall + back toward 1-col so full-tensor or 1-col H-strip can still succeed. + """ + n = _resolve_num_columns( + requested, + self.out_channels, + self.in_channels, + self.groups, + self.is_depthwise, + max_cols=max_cols, + ) + while n > 1: + try: + self._validate_l1_fit(n) + return n + except AIEOperatorConstraintError: + n = _resolve_num_columns( + n - 1, + self.out_channels, + self.in_channels, + self.groups, + self.is_depthwise, + max_cols=max_cols, + ) + return n + + def _halo_plan(self, num_columns: Optional[int] = None): + """Return design ``_plan_halo_h_strip`` result when k>1 H-strip is active. + + None when full-input L1 fits, or config is not standard k>1, + or no DMA-legal L1 plan exists (including optional bottom/right extra pad). + Multi-col non-depthwise grouped uses group-block split (no H-strip). + """ + # Depthwise uses channel tiles; pointwise has its own H-strip path. + if self.is_depthwise or self._is_pointwise(): + return None + n = 1 + cols = max( + 1, + int(num_columns if num_columns is not None else self.effective_num_columns), + ) + # Grouped multi-col: design uses group-block split, not H-strip. + if self.groups > 1 and cols > 1: + return None + if self.groups == 1: + oc_per_col = self.out_channels // cols + else: + cols = 1 + oc_per_col = self.out_channels + in_spatial = self.in_height * self.in_width + out_spatial = self.out_height * self.out_width + weight_per_oc = ( + (self.in_channels // self.groups) + * self.kernel_size[0] + * self.kernel_size[1] + ) + weight_store_per_oc = weight_per_oc + _bias_per_oc(self.use_bias) + input_size = n * self.in_channels * in_spatial + budget = _L1_TRIPLE_BUDGET_BYTES + if oc_per_col <= 0: + return None + if self.groups == 1: + oc_tile = _choose_oc_tile( + oc_per_col, input_size, weight_store_per_oc, out_spatial, budget + ) + if _l1_triple_fits( + input_size, + oc_tile * weight_store_per_oc, + n * oc_tile * out_spatial, + budget, + ): + return None + else: + weight_size = self.out_channels * weight_store_per_oc + output_size = n * self.out_channels * out_spatial + if _l1_triple_fits(input_size, weight_size, output_size, budget): + return None + ph, pw = self.padding + kh, kw = self.kernel_size + sh, sw = self.stride + return _plan_halo_h_strip( + self.in_height, + self.in_width, + self.out_height, + self.out_width, + self.in_channels, + oc_per_col, + weight_store_per_oc, + kh, + kw, + sh, + sw, + ph, + pw, + budget, + ) + + def _uses_halo_spatial_tiling(self, num_columns: Optional[int] = None) -> bool: + """True when design enables k>1 host-pad H-strip (groups==1).""" + return self._halo_plan(num_columns) is not None + + def _host_pad_input_nchw( + self, x_nchw: torch.Tensor, plan: Optional[dict] = None + ) -> torch.Tensor: + """Zero-pad (C,H,W) for k>1 spatial design (conv pad + optional DMA pad). + + Conv padding is applied symmetrically; any DMA extra is **bottom/right** + only so true top-left outputs match the unpadded formula. + """ + ph, pw = self.padding + extra_h = 0 + extra_w = 0 + if plan is not None: + extra_h = int(plan.get("extra_h", 0)) + extra_w = int(plan.get("extra_w", 0)) + if ph == 0 and pw == 0 and extra_h == 0 and extra_w == 0: + return x_nchw.contiguous() + # F.pad pad order: (W_left, W_right, H_top, H_bottom) + return torch.nn.functional.pad( + x_nchw, (pw, pw + extra_w, ph, ph + extra_h) + ).contiguous() + + def _pad_input_xrt(self, in_b: XRTTensor) -> XRTTensor: + """Pad host/runtime input buffer when k>1 spatial L3 expects padded size.""" + plan = self._halo_plan() + if plan is None: + return in_b + t = in_b.to_torch() + if not isinstance(t, torch.Tensor): + t = torch.tensor(t) + t = t.detach().cpu().contiguous() + if t.dtype != torch.bfloat16: + t = t.to(torch.bfloat16) + flat = t.reshape(-1) + expect = self.in_channels * self.in_height * self.in_width + padded_n = self.in_channels * plan["padded_h"] * plan["padded_w"] + if flat.numel() == padded_n: + return in_b + if flat.numel() != expect: + raise AIEOperatorConstraintError( + f"AIEConv2d halo-spatial pad expected {expect} elems " + f"(or already-padded {padded_n}), got {flat.numel()}" + ) + x_nchw = flat.reshape(self.in_channels, self.in_height, self.in_width) + x_pad = self._host_pad_input_nchw(x_nchw, plan).reshape(-1).contiguous() + return XRTTensor.from_torch(x_pad) + + def _crop_npu_output_to_true( + self, npu_out: XRTTensor, true_out: XRTTensor, plan: dict + ) -> None: + """Copy design-spatial NPU out into true OH×OW host out buffer.""" + design_oh = plan["design_oh"] + design_ow = plan["design_ow"] + true_oh = self.out_height + true_ow = self.out_width + t = npu_out.to_torch() + if not isinstance(t, torch.Tensor): + t = torch.tensor(t) + t = t.detach().cpu().contiguous() + if t.dtype != torch.bfloat16: + t = t.to(torch.bfloat16) + vol = t.reshape(self.out_channels, design_oh, design_ow) + cropped = vol[:, :true_oh, :true_ow].contiguous().reshape(-1) + if cropped.dtype == torch.bfloat16: + np_c = cropped.view(torch.uint16).numpy().view(np.dtype("bfloat16")) + else: + np_c = cropped.numpy().astype(bfloat16, copy=False) + true_out.data.reshape(-1)[:] = np_c + if hasattr(true_out, "_sync_to_device"): + true_out._sync_to_device() + + def _validate_l1_fit(self, num_columns: int) -> None: + """Raise if the design's L1 triple (in+weight+out, bf16) cannot fit. + + Mirrors design.py tile selection: groups==1 OC-tiles with + full input in L1, or H-strip spatial (pointwise or k>1 host-pad RF) + when full input exceeds budget; depthwise channel-tiles; non-DW + groups multi-col uses per-col IC/OC group blocks (must fit L1); + 1-col groups may use full triple or H-strip. Multi-column OC split + does not reduce full-input L1 for groups==1 (input broadcast). + """ + n = 1 # MLIR is specialized for N=1; batch is looped on host. + in_spatial = self.in_height * self.in_width + out_spatial = self.out_height * self.out_width + weight_per_oc = ( + (self.in_channels // self.groups) + * self.kernel_size[0] + * self.kernel_size[1] + ) + # L1 weight footprint includes packed bias (+1 per OC/channel) when used. + weight_store_per_oc = weight_per_oc + _bias_per_oc(self.use_bias) + input_size = n * self.in_channels * in_spatial + budget = _L1_TRIPLE_BUDGET_BYTES + bpe = _BYTES_PER_BF16 + cols = max(1, int(num_columns)) + is_pointwise = self._is_pointwise() + + if self.is_depthwise: + c_per_col = self.in_channels // cols + c_tile = _choose_channel_tile( + c_per_col, in_spatial, out_spatial, weight_store_per_oc, budget + ) + tile_elems = c_tile * (in_spatial + weight_store_per_oc + out_spatial) + if tile_elems * bpe > budget: + need = tile_elems * bpe + raise AIEOperatorConstraintError( + f"AIEConv2d depthwise L1 footprint exceeds budget: " + f"min channel tile needs ~{need} bytes " + f"(budget {_L1_TRIPLE_BUDGET_BYTES} bytes for in+weight+out " + f"bf16 at depth=1). Config: C={self.in_channels}, " + f"spatial={self.in_height}x{self.in_width}→" + f"{self.out_height}x{self.out_width}, " + f"kernel={self.kernel_size}, cols={cols}. " + f"Reduce spatial size/channels or wait for spatial L1 tiling." + ) + return + + if self.groups == 1: + oc_per_col = self.out_channels // cols + oc_tile = _choose_oc_tile( + oc_per_col, input_size, weight_store_per_oc, out_spatial, budget + ) + full_fits = _l1_triple_fits( + input_size, + oc_tile * weight_store_per_oc, + n * oc_tile * out_spatial, + budget, + ) + if full_fits: + return + + # Pointwise H-strip can still fit when full input does not. + # Prefer full oc_per_col per strip (num_oc_tiles=1; no DMA stride-0). + if is_pointwise: + tile_h = _choose_h_tile_pointwise( + self.in_height, + self.in_channels, + self.in_width, + oc_per_col, + weight_store_per_oc, + budget, + ) + in_tile = n * self.in_channels * tile_h * self.in_width + out_tile_sp = tile_h * self.out_width + if _l1_triple_fits( + in_tile, + oc_per_col * weight_store_per_oc, + n * oc_per_col * out_tile_sp, + budget, + ): + return + need = ( + in_tile + + oc_per_col * weight_store_per_oc + + n * oc_per_col * out_tile_sp + ) * bpe + raise AIEOperatorConstraintError( + f"AIEConv2d pointwise L1 footprint exceeds budget even with " + f"H-strip spatial tiling (full OC/col): needs ~{need} bytes " + f"(budget {_L1_TRIPLE_BUDGET_BYTES}; tile_h={tile_h}). " + f"Config: IC={self.in_channels}, OC={self.out_channels}, " + f"spatial={self.in_height}x{self.in_width}, cols={cols}." + ) + + # k>1 host-pad RF H-strip with full oc_per_col (+ DMA pad). + if self._halo_plan(cols) is not None: + return + + ph, pw = self.padding + padded_w = self.in_width + 2 * pw + kh, sh = self.kernel_size[0], self.stride[0] + tile_oh = _choose_h_tile_standard( + self.out_height, + self.in_channels, + padded_w, + oc_per_col, + weight_store_per_oc, + self.out_width, + kh, + sh, + budget, + ) + in_h_tile = _rf_in_h(max(1, tile_oh), sh, kh) + in_tile = n * self.in_channels * in_h_tile * padded_w + out_tile_sp = max(1, tile_oh) * self.out_width + need = ( + in_tile + + oc_per_col * weight_store_per_oc + + n * oc_per_col * out_tile_sp + ) * bpe + input_bytes = input_size * bpe + # Distinguish true L1 OOM from DMA-parity impossibility. + out_strip = max(1, tile_oh) * self.out_width + in_strip = in_h_tile * padded_w + dma_ok = (out_strip % 2 == 0) and (in_strip % 2 == 0) + dma_note = "" + if need <= budget and not dma_ok: + dma_note = ( + f" Natural strip sizes are not DMA-aligned " + f"(out_strip={out_strip}, in_strip={in_strip} elems) and no " + f"bottom/right DMA extra-pad plan found within search bound." + ) + raise AIEOperatorConstraintError( + f"AIEConv2d L1 footprint exceeds budget even with k>1 " + f"host-pad H-strip spatial tiling: needs ~{need} bytes " + f"(budget {_L1_TRIPLE_BUDGET_BYTES} bytes; " + f"full input alone is {input_bytes} bytes; " + f"tile_oh={tile_oh}). " + f"Config: IC={self.in_channels}, OC={self.out_channels}, " + f"spatial={self.in_height}x{self.in_width}→" + f"{self.out_height}x{self.out_width}, " + f"kernel={self.kernel_size}, cols={cols}. " + f"Note: multi-column OC split does not reduce input L1 " + f"(input is broadcast per column).{dma_note}" + ) + + # Non-depthwise grouped: multi-col group-block (per-col IC/OC) or + # 1-col full tensor / k>1 host-pad H-strip. + if self.groups % cols != 0: + raise AIEOperatorConstraintError( + f"AIEConv2d grouped multi-col requires groups % cols == 0, " + f"got groups={self.groups}, cols={cols}" + ) + ic_per_col = self.in_channels // cols + oc_per_col = self.out_channels // cols + in_col = n * ic_per_col * in_spatial + w_col = oc_per_col * weight_store_per_oc + out_col = n * oc_per_col * out_spatial + if _l1_triple_fits(in_col, w_col, out_col, budget): + return + if cols > 1: + # Multi-col grouped has no H-strip path; caller may drop columns. + need = (in_col + w_col + out_col) * bpe + raise AIEOperatorConstraintError( + f"AIEConv2d grouped multi-col (groups={self.groups}, cols={cols}) " + f"per-col L1 triple needs ~{need} bytes " + f"(budget {_L1_TRIPLE_BUDGET_BYTES}). " + f"Config: IC={self.in_channels}, OC={self.out_channels}, " + f"spatial={self.in_height}x{self.in_width}. " + f"Try fewer columns or smaller spatial (1-col may H-strip)." + ) + if self._halo_plan(1) is not None: + return + weight_size = self.out_channels * weight_store_per_oc + output_size = n * self.out_channels * out_spatial + triple = (input_size + weight_size + output_size) * bpe + raise AIEOperatorConstraintError( + f"AIEConv2d grouped (groups={self.groups}, non-depthwise) " + f"requires full in+weight+out in L1 (~{triple} bytes) or a legal " + f"k>1 H-strip plan, but budget is {_L1_TRIPLE_BUDGET_BYTES} bytes. " + f"Config: IC={self.in_channels}, OC={self.out_channels}, " + f"spatial={self.in_height}x{self.in_width}." + ) + + def set_up_artifacts(self): + """Set up compilation artifacts (L1 tiles + multi-col split).""" + operator_dir = Path(__file__).parent + design_path = operator_dir / "design.py" + + try: + dev = aie_utils.get_current_device() + kernel_dir = "aie2p" if getattr(dev, "cols", 4) > 4 else "aie2" + except Exception: + kernel_dir = "aie2" + dev = None + + if dev is None: + try: + dev = aie_utils.get_current_device() + except Exception: + from aie.iron.device import NPU1 + + dev = NPU1() + + # Column cap from target device model (NPU1.cols / NPU2.cols). + max_cols = getattr(dev, "cols", None) or 4 + effective_num_columns = self._resolve_columns_for_l1( + self.requested_num_columns, max_cols=max_cols + ) + self.effective_num_columns = effective_num_columns + # L1 re-check after device/group/L1 column clamp. + self._validate_l1_fit(effective_num_columns) + + bias_tag = "bias" if self.use_bias else "nobias" + file_name_base = ( + f"conv2d_{self.in_channels}_{self.out_channels}_{self.in_height}x{self.in_width}_" + f"{self.kernel_size[0]}x{self.kernel_size[1]}_" + f"s{self.stride[0]}x{self.stride[1]}_" + f"p{self.padding[0]}x{self.padding[1]}_" + f"g{self.groups}_{effective_num_columns}c_{bias_tag}" + ) + + mlir_artifact = PythonGeneratedMLIRArtifact( + f"{file_name_base}.mlir", + DesignGenerator( + design_path, + "my_conv2d", + args=(), + kwargs={ + "dev": dev, + "N": 1, + "in_channels": self.in_channels, + "in_height": self.in_height, + "in_width": self.in_width, + "out_channels": self.out_channels, + "out_height": self.out_height, + "out_width": self.out_width, + "kernel_h": self.kernel_size[0], + "kernel_w": self.kernel_size[1], + "stride_h": self.stride[0], + "stride_w": self.stride[1], + "pad_h": self.padding[0], + "pad_w": self.padding[1], + "groups": self.groups, + "use_bias": self.use_bias, + "num_columns": effective_num_columns, + "tile_size": self.tile_size, + "trace_size": 0, + }, + ), + ) + + kernel_obj = KernelObjectArtifact( + "conv2d.o", + dependencies=[ + SourceArtifact( + self.context.base_dir / "aie_kernels" / kernel_dir / "conv2d.cc" + ) + ], + ) + + xclbin_artifact = XclbinArtifact( + f"{file_name_base}.xclbin", + mlir_input=mlir_artifact, + dependencies=[mlir_artifact, kernel_obj], + extra_flags=[], + ) + + insts_artifact = InstsBinArtifact( + f"{file_name_base}.bin", + mlir_input=mlir_artifact, + dependencies=[mlir_artifact], + ) + + self.xclbin_artifact = xclbin_artifact + self.insts_artifact = insts_artifact + + self.add_artifacts([xclbin_artifact, insts_artifact]) + + def compile(self, dry_run: bool = False): + """Compile artifacts; invalidate cached NPU callable.""" + result = super().compile(dry_run=dry_run) + self._callable = None + return result + + def _get_op_callable(self) -> Callable[..., Any]: + """Lazy get_callable after compile (maxpool-style cache).""" + if self._callable is None: + if not self.artifacts: + self.compile() + self._callable = self.get_callable() + return self._callable + + def __call__( + self, + x: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + return self.forward(x, weight, bias) + + def forward( + self, + x: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ): + """Run conv2d (compile if needed; pack bias into weight DMA when used).""" + if len(x.shape) != 4: + raise AIEOperatorConstraintError( + f"AIEConv2d expects 4D input (N, C, H, W), got shape {x.shape}" + ) + + batch_size, actual_in_channels, actual_in_height, actual_in_width = x.shape + + if actual_in_channels != self.in_channels: + raise AIEOperatorConstraintError( + f"Expected {self.in_channels} input channels, got {actual_in_channels}" + ) + if actual_in_height != self.in_height or actual_in_width != self.in_width: + raise AIEOperatorConstraintError( + f"AIEConv2d configured for HxW=({self.in_height},{self.in_width}), " + f"but got input spatial {actual_in_height}x{actual_in_width} (shape {x.shape})" + ) + + outputs = [] + for n in range(batch_size): + x_n = x[n].contiguous() + result_n = self._process_single(x_n, weight, bias) + outputs.append(result_n) + + return torch.stack(outputs, dim=0) + + def _process_single( + self, + x: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ): + """Process a single sample (C, H, W) via NPU + optional host bias.""" + x_flat = x.reshape(-1).contiguous() + if x_flat.dtype != torch.bfloat16: + x_flat = x_flat.to(torch.bfloat16) + + weight_flat = weight.reshape(-1).contiguous() + if weight_flat.dtype != torch.bfloat16: + weight_flat = weight_flat.to(torch.bfloat16) + + if x_flat.numel() != self.input_size: + raise AIEOperatorConstraintError( + f"Flattened input size {x_flat.numel()} != configured {self.input_size}" + ) + if weight_flat.numel() != self.weight_size: + raise AIEOperatorConstraintError( + f"Flattened weight size {weight_flat.numel()} != configured {self.weight_size}" + ) + + op_func = self._get_op_callable() + in_b = XRTTensor.from_torch(x_flat) + w_b = XRTTensor.from_torch(weight_flat) + out_b = XRTTensor((self.output_size,), dtype=bfloat16) + + if self.use_bias and self.bias_size > 0: + # get_callable expects 4 args when use_bias; zeros if bias omitted. + if bias is None: + bias_t = torch.zeros(self.bias_size, dtype=torch.bfloat16) + else: + bias_t = bias.contiguous() + if bias_t.dtype != torch.bfloat16: + bias_t = bias_t.to(torch.bfloat16) + bias_b = XRTTensor.from_torch(bias_t) + op_func(in_b, w_b, bias_b, out_b) + else: + op_func(in_b, w_b, out_b) + + # Clone off XRT BO before buffers leave scope (batch>1 stack safety). + result = out_b.to_torch() + if not isinstance(result, torch.Tensor): + result = torch.tensor(result) + if result.dtype != torch.bfloat16: + result = result.to(torch.bfloat16) + result = result.detach().cpu().contiguous().clone() + + return result.reshape(self.out_channels, self.out_height, self.out_width) + + def _host_apply_bias(self, out_buf, bias_buf) -> None: + """In-place host bias add on XRT output buffer (bf16). + + Uses to_torch() so any device→host sync performed by the runtime is + honored, then writes the summed result back through the mapped ``data`` + view (verified writable for XRTTensor). + """ + out_t = out_buf.to_torch().reshape( + self.out_channels, self.out_height, self.out_width + ) + bias_t = ( + bias_buf.to_torch().to(dtype=out_t.dtype).reshape(self.out_channels, 1, 1) + ) + summed = (out_t + bias_t).contiguous().reshape(-1) + # Convert torch bf16 → numpy bf16 without float32 round-trip when possible. + if summed.dtype == torch.bfloat16: + np_sum = ( + summed.detach() + .cpu() + .view(torch.uint16) + .numpy() + .view(np.dtype("bfloat16")) + ) + else: + np_sum = summed.detach().cpu().numpy().astype(bfloat16, copy=False) + out_buf.data.reshape(-1)[:] = np_sum + # Critical: to_torch()/numpy() sync FROM device and would wipe host + # writes unless we push the biased result back to the device BO. + if hasattr(out_buf, "_sync_to_device"): + out_buf._sync_to_device() + + def get_arg_spec(self): + """Runtime arg specs for run_test / high-level path. + + Host-facing order: + - with bias: in, weight, bias, out (bias applied on host after NPU) + - without: in, weight, out + + NPU instruction sequence is always (in, weight, out); get_callable + strips the bias buffer before DefaultNPURuntime.run. + """ + # Sizes used by run_test buffer allocation / XRTTensor shapes. + input_size = self.in_channels * self.in_height * self.in_width + weight_size = ( + self.out_channels + * self.in_channels + // self.groups + * self.kernel_size[0] + * self.kernel_size[1] + ) + output_size = self.out_channels * self.out_height * self.out_width + # Cache for callers that read these attributes. + self.input_size = input_size + self.weight_size = weight_size + self.output_size = output_size + + specs = [ + AIERuntimeArgSpec("in", (input_size,)), + AIERuntimeArgSpec("in", (weight_size,)), + ] + if self.use_bias and self.bias_size > 0: + specs.append(AIERuntimeArgSpec("in", (self.bias_size,))) + specs.append(AIERuntimeArgSpec("out", (output_size,))) + return specs + + def _design_tile_channels(self) -> int: + """OC/channel tile size matching design.py L1 selection (for bias pack).""" + n = 1 + cols = max(1, int(self.effective_num_columns)) + in_spatial = self.in_height * self.in_width + out_spatial = self.out_height * self.out_width + weight_per_oc = ( + (self.in_channels // self.groups) + * self.kernel_size[0] + * self.kernel_size[1] + ) + w_store = weight_per_oc + _bias_per_oc(self.use_bias) + budget = _L1_TRIPLE_BUDGET_BYTES + if self.is_depthwise: + c_per_col = self.in_channels // cols + c_tile = _choose_channel_tile( + c_per_col, in_spatial, out_spatial, w_store, budget + ) + if c_per_col % c_tile != 0: + c_tile = c_per_col + return c_tile + if self.groups != 1: + # Grouped multi-col: full OC block per column (num_tiles=1). + # 1-col H-strip also uses full out_channels as the weight tile. + return self.out_channels // cols + oc_per_col = self.out_channels // cols + input_size = n * self.in_channels * in_spatial + oc_tile = _choose_oc_tile(oc_per_col, input_size, w_store, out_spatial, budget) + if oc_per_col % oc_tile != 0: + oc_tile = oc_per_col + full_fits = _l1_triple_fits( + input_size, oc_tile * w_store, n * oc_tile * out_spatial, budget + ) + if full_fits: + return oc_tile + # H-strip paths prefer full oc_per_col (num_oc_tiles==1). + if self._is_pointwise() or self._halo_plan(cols) is not None: + return oc_per_col + return oc_tile + + def _pack_weight_bias_xrt(self, w_buf, bias_buf) -> XRTTensor: + """Build L3 packed ``[W_tile‖B_tile]…`` tensor for on-device bias.""" + w_t = w_buf.to_torch().reshape(-1).contiguous() + b_t = bias_buf.to_torch().reshape(-1).contiguous() + if w_t.dtype != torch.bfloat16: + w_t = w_t.to(torch.bfloat16) + if b_t.dtype != torch.bfloat16: + b_t = b_t.to(torch.bfloat16) + # pack_weights_with_bias is numpy-oriented; convert via uint16 view. + w_np = w_t.detach().cpu().view(torch.uint16).numpy().view(np.dtype("bfloat16")) + b_np = b_t.detach().cpu().view(torch.uint16).numpy().view(np.dtype("bfloat16")) + packed = pack_weights_with_bias( + w_np, + b_np, + out_channels=self.out_channels, + in_channels=self.in_channels, + groups=self.groups, + kernel_h=self.kernel_size[0], + kernel_w=self.kernel_size[1], + num_columns=self.effective_num_columns, + is_depthwise=self.is_depthwise, + tile_channels=self._design_tile_channels(), + ) + packed_u16 = packed.view(np.uint16) + packed_t = torch.from_numpy(packed_u16.copy()).view(torch.bfloat16) + return XRTTensor.from_torch(packed_t.contiguous()) + + def get_callable(self): + """Callable that packs bias into weights and runs NPU conv (≤2 DMAs).""" + if self.xclbin_artifact is None or self.insts_artifact is None: + self.set_up_artifacts() + npu_kernel = NPUKernel( + xclbin_path=self.xclbin_artifact.filename, + kernel_name=self.xclbin_artifact.kernel_name, + insts_path=self.insts_artifact.filename, + ) + handle = aie_utils.DefaultNPURuntime.load(npu_kernel) + use_bias = self.use_bias and self.bias_size > 0 + + def call(*args): + # k>1 spatial designs use host-padded L3 input (design input_ty). + # External API / run_test still pass unpadded C*H*W; pad here. + # When DMA pad grows design OH/OW, stage a larger NPU out and crop. + plan = self._halo_plan() + need_stage_out = False + design_out_size = 0 + if plan is not None: + design_out_size = ( + self.out_channels * plan["design_oh"] * plan["design_ow"] + ) + need_stage_out = design_out_size > ( + self.out_channels * self.out_height * self.out_width + ) + + def _run_npu(in_buf, w_buf, out_buf): + in_p = self._pad_input_xrt(in_buf) + if need_stage_out: + npu_out = XRTTensor((design_out_size,), dtype=bfloat16) + result = aie_utils.DefaultNPURuntime.run( + handle, [in_p, w_buf, npu_out] + ) + self._crop_npu_output_to_true(npu_out, out_buf, plan) + return result + return aie_utils.DefaultNPURuntime.run(handle, [in_p, w_buf, out_buf]) + + if use_bias: + if len(args) != 4: + raise ValueError( + f"AIEConv2d with bias expects 4 args (in, weight, bias, out), got {len(args)}" + ) + in_b, w_b, bias_b, out_b = args + packed_w = self._pack_weight_bias_xrt(w_b, bias_b) + return _run_npu(in_b, packed_w, out_b) + if len(args) < 3: + raise ValueError( + f"AIEConv2d expects (in, weight, out), got {len(args)} args" + ) + return _run_npu(args[0], args[1], args[2]) + + return call diff --git a/iron/operators/conv2d/reference.py b/iron/operators/conv2d/reference.py new file mode 100644 index 00000000..b0e7c989 --- /dev/null +++ b/iron/operators/conv2d/reference.py @@ -0,0 +1,217 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU golden for AIEConv2d tests via torch F.conv2d. + +``conv2d_cpu`` wraps F.conv2d; ``generate_golden_reference`` builds seeded +tensors and expected output. HW tests use bf16 tolerances (tolerances.py). +""" + +import torch +import torch.nn.functional as F +from typing import Tuple, Union + + +def conv2d_cpu( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor = None, + stride: Union[int, Tuple[int, int]] = 1, + padding: Union[int, Tuple[int, int]] = 0, + dilation: Union[int, Tuple[int, int]] = 1, + groups: int = 1, +) -> torch.Tensor: + """F.conv2d wrapper used as the golden for all conv2d tests.""" + output = F.conv2d( + input=input, + weight=weight, + bias=bias, + stride=stride, + padding=padding, + dilation=dilation, + groups=groups, + ) + return output + + +def generate_golden_reference( + batch_size: int = 1, + in_channels: int = 3, + in_height: int = 32, + in_width: int = 32, + out_channels: int = 16, + kernel_size: Union[int, Tuple[int, int]] = 3, + stride: Union[int, Tuple[int, int]] = 1, + padding: Union[int, Tuple[int, int]] = 0, + dilation: Union[int, Tuple[int, int]] = 1, + groups: int = 1, + use_bias: bool = True, + dtype: torch.dtype = torch.bfloat16, + seed: int = 42, +): + """Seeded tensors + expected output via conv2d_cpu (bf16 drawn in fp32 then cast).""" + torch.manual_seed(seed) + + # Normalize kernel_size, stride, padding, dilation to tuples + if isinstance(kernel_size, int): + kernel_size = (kernel_size, kernel_size) + if isinstance(stride, int): + stride = (stride, stride) + if isinstance(padding, int): + padding = (padding, padding) + if isinstance(dilation, int): + dilation = (dilation, dilation) + + # Validate groups + assert in_channels % groups == 0, "in_channels must be divisible by groups" + assert out_channels % groups == 0, "out_channels must be divisible by groups" + + out_height = calculate_output_dim( + in_height, kernel_size[0], stride[0], padding[0], dilation[0] + ) + out_width = calculate_output_dim( + in_width, kernel_size[1], stride[1], padding[1], dilation[1] + ) + + if dtype == torch.bfloat16: + input_tensor = ( + torch.randn( + batch_size, in_channels, in_height, in_width, dtype=torch.float32 + ) + * 2.0 + ) + input_tensor = input_tensor.to(dtype) + else: + input_tensor = ( + torch.randn(batch_size, in_channels, in_height, in_width, dtype=dtype) * 2.0 + ) + + # Create weight tensor + weight_shape = (out_channels, in_channels // groups, kernel_size[0], kernel_size[1]) + if dtype == torch.bfloat16: + weight_tensor = torch.randn(weight_shape, dtype=torch.float32) * 2.0 + weight_tensor = weight_tensor.to(dtype) + else: + weight_tensor = torch.randn(weight_shape, dtype=dtype) * 2.0 + + # Create bias tensor (if used) + bias_tensor = None + if use_bias: + if dtype == torch.bfloat16: + bias_tensor = torch.randn(out_channels, dtype=torch.float32) * 2.0 + bias_tensor = bias_tensor.to(dtype) + else: + bias_tensor = torch.randn(out_channels, dtype=dtype) * 2.0 + + expected_output = conv2d_cpu( + input=input_tensor, + weight=weight_tensor, + bias=bias_tensor, + stride=stride, + padding=padding, + dilation=dilation, + groups=groups, + ) + + assert ( + expected_output.shape[2] == out_height and expected_output.shape[3] == out_width + ), ( + f"Output shape mismatch in golden ref: F.conv2d gave {expected_output.shape[2:]} " + f"but formula gave ({out_height}, {out_width})" + ) + + return { + "input": input_tensor, + "weight": weight_tensor, + "bias": bias_tensor, + "output": expected_output, + "config": { + "batch_size": batch_size, + "in_channels": in_channels, + "in_height": in_height, + "in_width": in_width, + "out_channels": out_channels, + "kernel_size": kernel_size, + "stride": stride, + "padding": padding, + "dilation": dilation, + "groups": groups, + "use_bias": use_bias, + "out_height": out_height, + "out_width": out_width, + }, + } + + +def calculate_output_dim( + input_dim: int, + kernel_dim: int, + stride: int, + padding: int, + dilation: int, +) -> int: + """floor((input + 2*pad - dilation*(kernel-1) - 1) / stride + 1).""" + return (input_dim + 2 * padding - dilation * (kernel_dim - 1) - 1) // stride + 1 + + +if __name__ == "__main__": + # Quick test with simple configuration + print("Testing Conv2D CPU Reference Implementation...") + + # Test 1: Basic 3x3 convolution + golden = generate_golden_reference( + batch_size=1, + in_channels=3, + in_height=32, + in_width=32, + out_channels=16, + kernel_size=3, + stride=1, + padding=1, + groups=1, + ) + + print(f"\nTest 1: Basic 3x3 Conv") + print(f" Input shape: {golden['input'].shape}") + print(f" Weight shape: {golden['weight'].shape}") + print(f" Output shape: {golden['output'].shape}") + print(f" Config: {golden['config']}") + + # Test 2: Depthwise convolution + golden_dw = generate_golden_reference( + batch_size=1, + in_channels=16, + in_height=32, + in_width=32, + out_channels=16, + kernel_size=3, + stride=1, + padding=1, + groups=16, # Depthwise + ) + + print(f"\nTest 2: Depthwise 3x3 Conv") + print(f" Input shape: {golden_dw['input'].shape}") + print(f" Weight shape: {golden_dw['weight'].shape}") + print(f" Output shape: {golden_dw['output'].shape}") + print(f" Groups: {golden_dw['config']['groups']}") + + # Test 3: Strided convolution + golden_stride = generate_golden_reference( + batch_size=1, + in_channels=3, + in_height=64, + in_width=64, + out_channels=32, + kernel_size=3, + stride=2, + padding=1, + groups=1, + ) + + print(f"\nTest 3: Strided 3x3 Conv (stride=2)") + print(f" Input shape: {golden_stride['input'].shape}") + print(f" Output shape: {golden_stride['output'].shape}") + print(f" Config: {golden_stride['config']}") + + print("\nAll tests passed!") diff --git a/iron/operators/conv2d/test.py b/iron/operators/conv2d/test.py new file mode 100644 index 00000000..7fc3386b --- /dev/null +++ b/iron/operators/conv2d/test.py @@ -0,0 +1,554 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""NPU end-to-end tests for AIEConv2d (CPU-only coverage in cpu_test.py).""" + +import pytest + +import torch + +from iron.operators.conv2d.op import AIEConv2d +from iron.operators.conv2d.reference import ( + generate_golden_reference, + calculate_output_dim, +) +from iron.operators.conv2d.tolerances import hw_tolerances +from iron.common import AIEOperatorConstraintError +from iron.common.test_utils import run_test + + +def get_params(): + """Generate conv2d pytest params (device cols, shapes, extensive marks). + + Defensive device query so collect-only / CPU-only envs do not crash. + """ + import aie.utils as aie_utils + + # Defensive device discovery (pure-CPU reference tests + collectonly safety) + max_cols = 4 + try: + dev = aie_utils.get_current_device() + max_cols = dev.cols + except Exception: + pass + + # Core configurations (in_ch, out_ch, k, s, p, g, use_bias) + # Extended set for good coverage of variants (exercises all golden paths, + # column chunking, bias ObjectFifo singular broadcast, variant kernels, + # conditional rt.sequence, and prepare_runtime runlist arity). + configs = [ + (3, 16, 3, 1, 1, 1, True), # basic +bias + (3, 16, 3, 1, 1, 1, False), # basic nobias + (16, 16, 3, 1, 1, 1, True), + (16, 16, 3, 1, 1, 16, True), # depthwise +bias + (16, 16, 3, 1, 1, 16, False), # depthwise nobias + (32, 64, 1, 1, 0, 1, True), # pointwise + (32, 64, 1, 1, 0, 1, False), + (16, 32, 3, 2, 1, 1, True), # strided +pad + (16, 32, 3, 2, 0, 1, True), # strided no pad + (8, 16, 3, 1, 2, 2, True), # groups=2 + (4, 8, 3, 1, 1, 2, True), + ] + + # Explicit core configs for regular marking (robust vs list order / slicing). + # Regular CI coverage: + # - 3→16 bias/nobias: host-bias + full/near-full L1 + # - 16→16 groups=1 bias: multi-tile OC path at 32x32 (oc_tile=8) + # - 16 depthwise bias: multi-tile channel path at 32x32 (c_tile=8) + # Also 16x16 CORE @ 2c (OC/channel split, ≤2 DMA, host bias). + CORE_CONFIGS = [ + (3, 16, 3, 1, 1, 1, True), + (3, 16, 3, 1, 1, 1, False), + (16, 16, 3, 1, 1, 1, True), # standard multi-tile OC + (16, 16, 3, 1, 1, 16, True), # depthwise multi-tile channels + ] + + # 16x16 + 32x32 CORE @ 1c: L1 fit. 16x16 CORE @ 2c: multi-col smoke. + # 32x32+ multi-col and 64 spatial stay extensive until proven green. + spatials = [(16, 16), (32, 32), (64, 64)] + col_candidates = [1, 2, 4, 8] + + params = [] + for h, w in spatials: + for cfg in configs: + in_ch, out_ch, k, s, p, g, use_bias = cfg + for nc in col_candidates: + if nc > max_cols: + continue + + # Dilation is fixed to 1 in current AIEConv2d (asserted in op.py). + # Use the *shared* calculate_output_dim from reference (exact match + # to generate_golden_reference + operator + design for d=1). + # This guarantees the out_h/out_w used for divisibility + naming + # are identical to those in the golden "output" tensor shape. + dilation = 1 + out_h = calculate_output_dim(h, k, s, p, dilation) + out_w = calculate_output_dim(w, k, s, p, dilation) + + # Sizes that must be evenly divisible for column chunking on + # *flattened* elements (critical: design.py chunks C*H*W, weight, + # and output by num_aie_columns for parallel columns). + in_size = in_ch * h * w # N=1 (MLIR specialization) + w_size = out_ch * (in_ch // g) * k * k + out_size = out_ch * out_h * out_w + + if ( + nc == 0 + or in_size % nc != 0 + or w_size % nc != 0 + or out_size % nc != 0 + ): + continue + + tile_size = in_size // nc + + # Regular subset ("not extensive"): + # - 16x16 / 32x32 CORE @ 1c — L1 OC/channel tiles + # - 16x16 CORE @ 2c — multi-col OC/channel split smoke + # Bias remains host-side (2 input DMA limit per compute tile). + # Larger multi-col (4c/8c, 32x32+) stays extensive. + is_core_config = cfg in CORE_CONFIGS + is_regular = is_core_config and ( + (nc == 1 and (h, w) in ((16, 16), (32, 32))) + or (nc == 2 and (h, w) == (16, 16)) + ) + + marks = [] if is_regular else [pytest.mark.extensive] + + bias_str = "bias" if use_bias else "nobias" + name = f"conv2d_{in_ch}x{out_ch}_k{k}_s{s}_p{p}_g{g}_{bias_str}_{h}x{w}_{nc}c_{tile_size}t" + + # Note: batch always 1 for the low-level run_test path (N=1 MLIR specialization) + params.append( + pytest.param( + in_ch, + out_ch, + k, + s, + p, + g, + use_bias, + 1, + h, + w, + nc, + tile_size, + id=name, + marks=marks, + ) + ) + + return params + + +# get_params() (single source of truth) is invoked *directly* inside @parametrize +# Direct get_params() in @parametrize (no top-level all_params assignment). +# Called at collection time; safe due to defensive device query inside. + + +# Explicit constant for the parameter names used in @parametrize decorators. +# This is the production hardening (see conv3d) against "N names vs M values" +# collection crashes when get_params or FORWARD_CASES evolve. The order and +# count (12) must exactly match the 12-tuples yielded by get_params() and the +# pytest.param values in FORWARD_CASES. +CONV2D_TEST_PARAM_NAMES = ( + "in_channels,out_channels,kernel_size,stride,padding,groups," + "use_bias,batch,in_h,in_w,num_aie_columns,tile_size" +) + + +@pytest.mark.metrics( + Latency=r"Latency \(us\): (?P[\d\.]+)", + Bandwidth=r"Effective Bandwidth: (?P[\d\.e\+-]+) GB/s", +) +# Smoke path: mean NPU latency + BO-sum Effective BW. +@pytest.mark.parametrize( + CONV2D_TEST_PARAM_NAMES, + get_params(), +) +def test_conv2d( + in_channels, + out_channels, + kernel_size, + stride, + padding, + groups, + use_bias, + batch, + in_h, + in_w, + num_aie_columns, + tile_size, + aie_context, +): + """NPU end-to-end test via run_test (compile, run, verify, metrics prints).""" + golden_ref = generate_golden_reference( + batch_size=batch, + in_channels=in_channels, + in_height=in_h, + in_width=in_w, + out_channels=out_channels, + kernel_size=kernel_size, + stride=stride, + padding=padding, + groups=groups, + use_bias=use_bias, + seed=42, + ) + + # Construct may raise AIEOperatorConstraintError when no L1 plan fits. + try: + operator = AIEConv2d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + stride=stride, + padding=padding, + groups=groups, + use_bias=use_bias, + in_height=in_h, + in_width=in_w, + num_aie_columns=num_aie_columns, + tile_size=tile_size, + context=aie_context, + ) + except AIEOperatorConstraintError as e: + pytest.skip(f"Unsupported AIEConv2d config (L1/column constraint): {e}") + + # Cross-validate output dimension math (catches formula drift) + ref_out_shape = golden_ref["output"].shape + assert ref_out_shape[0] == batch + assert ref_out_shape[1] == out_channels + assert ( + operator.out_height == ref_out_shape[2] + ), f"out_height mismatch: operator={operator.out_height}, ref={ref_out_shape[2]}" + assert ( + operator.out_width == ref_out_shape[3] + ), f"out_width mismatch: operator={operator.out_width}, ref={ref_out_shape[3]}" + + # Prepare buffers (bias only when use_bias) + input_buffers = { + "input": golden_ref["input"], + "weight": golden_ref["weight"], + } + if use_bias and golden_ref["bias"] is not None: + input_buffers["bias"] = golden_ref["bias"] + + output_buffers = {"output": golden_ref["output"]} + + # Tolerances: see iron/operators/conv2d/tolerances.py (audit-backed HW policy). + tols = hw_tolerances() + errors, latency_us, bandwidth_gbps = run_test( + operator, + input_buffers, + output_buffers, + rel_tol=tols.rel_tol, + abs_tol=tols.abs_tol, + max_error_rate=tols.max_error_rate, + ) + + # Exactly the two lines required by the @metrics regexes (main-tree style, + # identical to maxpool/avgpool/conv3d/reduction). Extra debug prints removed + # for robust CSV/metrics reporter capture and pre-push hook compatibility. + print(f"\nLatency (us): {latency_us:.1f}") + print(f"Effective Bandwidth: {bandwidth_gbps:.6e} GB/s\n") + + assert not errors, f"Test failed with errors: {errors}" + + +# Carefully chosen representative cases for the high-level forward API test. +# Explicit pytest.param objects (maxpool/conv3d/avgpool/reduction pattern) guarantee: +# - Stable, descriptive test IDs for CSV/metrics and reports +# - No dependency on ordering/count of get_params() results (uses independent FORWARD_CASES) +# - No fragile slicing or mark introspection +# - Targeted coverage of column/tile variants (different MLIR specializations) +# - Bias on/off + key kernel variants (standard/depthwise/pointwise/strided) +# +# These deliberately stay small/fast even under --iterations while still +# exercising operator.compile() + forward()/__call__ (get_callable + XRTTensor) +# and the python-level batching over N=1-specialized MLIR. +FORWARD_CASES = [ + # 16x16 + 1-col keeps full tensors inside L1 (~64KB) with depth=1. + # tile_size = in_ch * H * W for nc=1. + pytest.param( + 3, + 16, + 3, + 1, + 1, + 1, + True, + 1, + 16, + 16, + 1, + 768, + id="conv2d_forward_basic_bias_16x16_1c", + ), + pytest.param( + 3, + 16, + 3, + 1, + 1, + 1, + False, + 1, + 16, + 16, + 1, + 768, + id="conv2d_forward_basic_nobias_16x16_1c", + ), + pytest.param( + 16, + 16, + 3, + 1, + 1, + 16, + True, + 1, + 16, + 16, + 1, + 4096, + id="conv2d_forward_depthwise_16x16_1c", + ), + pytest.param( + 8, + 16, + 1, + 1, + 0, + 1, + True, + 1, + 16, + 16, + 1, + 2048, + id="conv2d_forward_pointwise_16x16_1c", + ), + pytest.param( + 3, + 16, + 3, + 2, + 1, + 1, + True, + 1, + 16, + 16, + 1, + 768, + id="conv2d_forward_strided_16x16_1c", + ), +] + + +@pytest.mark.extensive +@pytest.mark.parametrize( + CONV2D_TEST_PARAM_NAMES, + FORWARD_CASES, +) +def test_conv2d_forward( + in_channels, + out_channels, + kernel_size, + stride, + padding, + groups, + use_bias, + batch, + in_h, + in_w, + num_aie_columns, + tile_size, + aie_context, +): + """Forward / __call__ path: compile, run, verify (incl. batch loop).""" + golden_ref = generate_golden_reference( + batch_size=batch, + in_channels=in_channels, + in_height=in_h, + in_width=in_w, + out_channels=out_channels, + kernel_size=kernel_size, + stride=stride, + padding=padding, + groups=groups, + use_bias=use_bias, + seed=42, + ) + + try: + operator = AIEConv2d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + stride=stride, + padding=padding, + groups=groups, + use_bias=use_bias, + in_height=in_h, + in_width=in_w, + num_aie_columns=num_aie_columns, + tile_size=tile_size, + context=aie_context, + ) + except AIEOperatorConstraintError as e: + pytest.skip(f"Unsupported AIEConv2d config (L1/column constraint): {e}") + + # Modern MLIROperator path (AIEContext no longer exposes compile_all / + # prepare_runtime). Matches maxpool/avgpool forward tests. + operator.compile() + + # N=1 forward via __call__ / forward (XRTTensor + get_callable + host bias) + result = operator( + golden_ref["input"], + golden_ref["weight"], + golden_ref["bias"], + ) + expected = golden_ref["output"] + + assert ( + result.shape == expected.shape + ), f"Shape mismatch: got {result.shape}, expected {expected.shape}" + + tols = hw_tolerances() + rel_tol, abs_tol = tols.rel_tol, tols.abs_tol + if not torch.allclose(result, expected, rtol=rel_tol, atol=abs_tol): + max_diff = (result - expected).abs().max().item() + pytest.fail(f"Results don't match. Max diff: {max_diff}") + + # Batch=2 reuse of already-prepared operator/runlist + # This validates: + # - N=1 MLIR specialization + Python batching wrapper produces correct + # per-sample results matching the full-batch golden from generate_... + # - Golden generation with batch_size=2 works identically (F.conv2d handles N). + golden_b2 = generate_golden_reference( + batch_size=2, + in_channels=in_channels, + in_height=in_h, + in_width=in_w, + out_channels=out_channels, + kernel_size=kernel_size, + stride=stride, + padding=padding, + groups=groups, + use_bias=use_bias, + seed=42, + ) + result_b2 = operator( + golden_b2["input"], + golden_b2["weight"], + golden_b2["bias"], + ) + expected_b2 = golden_b2["output"] + assert ( + result_b2.shape == expected_b2.shape + ), f"Batch-2 shape mismatch: got {result_b2.shape}, expected {expected_b2.shape}" + if not torch.allclose(result_b2, expected_b2, rtol=rel_tol, atol=abs_tol): + max_diff = (result_b2 - expected_b2).abs().max().item() + pytest.fail(f"Batch-2 results don't match. Max diff: {max_diff}") + + +# --------------------------------------------------------------------------- +# Frozen-shape multi-iter bench (Ring 1). Default smoke still uses mean @metrics. +# --------------------------------------------------------------------------- + + +@pytest.mark.extensive +@pytest.mark.parametrize("dummy", [pytest.param(None, id="multi_col_4c_8c_matrix")]) +def test_conv2d_multi_col_4c_8c_matrix_present(dummy): + """Extensive matrix includes 4c/8c where device + divisibility allow.""" + import aie.utils as aie_utils + + params = get_params() + # Collect nc from param ids / values (12-tuple: ... num_aie_columns is index 10) + cols_seen = {p.values[10] for p in params} + max_cols = 4 + try: + max_cols = aie_utils.get_current_device().cols + except Exception: + pass + assert 1 in cols_seen and 2 in cols_seen + if max_cols >= 4: + assert 4 in cols_seen, f"expected 4c cases in matrix, saw {sorted(cols_seen)}" + if max_cols >= 8: + assert 8 in cols_seen, f"expected 8c cases in matrix, saw {sorted(cols_seen)}" + multi = [p for p in params if p.values[10] >= 4] + assert len(multi) >= 4, f"too few multi-col>=4 cases: {len(multi)}" + + +@pytest.mark.parametrize( + "in_ch,out_ch,k,s,p,g,use_bias,h,w,nc", + [ + pytest.param(8, 16, 3, 1, 1, 2, True, 16, 16, 2, id="groups2_2c_bias_16x16"), + pytest.param(8, 16, 3, 1, 2, 2, False, 16, 16, 2, id="groups2_2c_nobias_pad2"), + pytest.param(4, 8, 3, 1, 1, 2, True, 16, 16, 2, id="groups2_small_2c_bias"), + ], +) +def test_conv2d_grouped_multicol_npu( + in_ch, out_ch, k, s, p, g, use_bias, h, w, nc, aie_context +): + """Non-depthwise grouped multi-col group-block split on NPU.""" + golden = generate_golden_reference( + batch_size=1, + in_channels=in_ch, + in_height=h, + in_width=w, + out_channels=out_ch, + kernel_size=k, + stride=s, + padding=p, + groups=g, + use_bias=use_bias, + seed=20260808, + ) + try: + op = AIEConv2d( + in_channels=in_ch, + out_channels=out_ch, + kernel_size=k, + stride=s, + padding=p, + groups=g, + use_bias=use_bias, + in_height=h, + in_width=w, + num_aie_columns=nc, + context=aie_context, + ) + except AIEOperatorConstraintError as e: + pytest.skip(f"Unsupported grouped multi-col config: {e}") + + assert ( + op.effective_num_columns == nc + ), f"expected effective cols={nc}, got {op.effective_num_columns}" + + input_buffers = { + "input": golden["input"], + "weight": golden["weight"], + } + if use_bias and golden["bias"] is not None: + input_buffers["bias"] = golden["bias"] + output_buffers = {"output": golden["output"]} + tols = hw_tolerances() + errors, latency_us, bandwidth_gbps = run_test( + op, + input_buffers, + output_buffers, + rel_tol=tols.rel_tol, + abs_tol=tols.abs_tol, + max_error_rate=tols.max_error_rate, + ) + print(f"\n[grouped multi-col] Latency (us): {latency_us:.1f}") + print(f"Effective Bandwidth: {bandwidth_gbps:.6e} GB/s\n") + assert not errors, f"grouped multi-col failed: {errors}" + + +# CPU reference: cpu_test.py. NPU smoke: pytest -m "not extensive". diff --git a/iron/operators/conv2d/tolerances.py b/iron/operators/conv2d/tolerances.py new file mode 100644 index 00000000..759d6237 --- /dev/null +++ b/iron/operators/conv2d/tolerances.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""NPU vs torch-bf16 golden tolerances for AIEConv2d. + +Defaults for ``verify_buffer`` on hardware paths. bf16 MAC order can differ +from ``F.conv2d(bf16)``; small/grouped shapes need a non-zero abs floor. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Final + + +@dataclass(frozen=True) +class Conv2dHWTolerances: + """Hardware verification policy for one compare surface.""" + + rel_tol: float + abs_tol: float + max_error_rate: float + + +# Default HW verify policy (NPU vs torch bf16 golden). +HW_DEFAULT: Final = Conv2dHWTolerances( + rel_tol=0.05, + abs_tol=0.5, + max_error_rate=0.02, +) + +# Looser alternate for experiments only (not default product path). +HW_LOOSE: Final = Conv2dHWTolerances( + rel_tol=0.1, + abs_tol=1.0, + max_error_rate=0.02, +) + +# Tighter alternate for large-channel / pointwise-only experiments. +HW_STRICT_POINTWISE: Final = Conv2dHWTolerances( + rel_tol=0.05, + abs_tol=0.25, + max_error_rate=0.0, +) + + +def hw_tolerances() -> Conv2dHWTolerances: + """Return the product default HW verification tolerances.""" + return HW_DEFAULT