[sparsepixels] support asymmetric pool sizes in sparsepixels sparse pooling

hls4ml/backends/vivado/passes/sparsepixels.py

@@ -62,9 +62,9 @@ def transform(self, model, node):
                 n.set_attr('hash_in_name', h_in)
                 n.set_attr('hash_out_name', h_out)
                 hash_map[name] = h_out
-                ps = n.get_attr('pool_size')
+                ph, pw = n.get_attr('pool_height'), n.get_attr('pool_width')
                 prev_h, prev_w = spatial.get(src, spatial.get(n.inputs[0], (0, 0)))
-                spatial[name] = (prev_h // ps, prev_w // ps)
+                spatial[name] = (prev_h // ph, prev_w // pw)
 
             elif isinstance(n, SparseFlatten):
                 src = n.inputs[0]
@@ -134,7 +134,10 @@ def transform(self, model, node):
 sparse_pooling2d_config = """struct config{index} {{
     static const unsigned n_sparse = {n_sparse};
     static const unsigned n_chan = {n_chan};
-    static const unsigned pool_size = {pool_size};
+    static const unsigned in_height = {in_height};
+    static const unsigned in_width = {in_width};
+    static const unsigned pool_height = {pool_height};
+    static const unsigned pool_width = {pool_width};
     typedef {accum_t.name} accum_t;
 }};\n"""
 
@@ -217,12 +220,14 @@ def format(self, node):
 # Average pooling takes an accum_t; max pooling does not. Both take the two parallelization factors.
 sparse_pooling2d_avg_call = (
     'sparse_pooling_avg<{input_t}, {output_t}, ap_uint<{hash_bits}>, {accum_t_name}, '
-    '{n_sparse}, {n_chan}, {pool_size}, {pixel_parallel_factor}, {chan_parallel_factor}>'
+    '{n_sparse}, {n_chan}, {in_height}, {in_width}, {pool_height}, {pool_width}, '
+    '{pixel_parallel_factor}, {chan_parallel_factor}>'
     '({input}, {output}, {hash_in}, {hash_out});'
 )
 sparse_pooling2d_max_call = (
     'sparse_pooling_max<{input_t}, {output_t}, ap_uint<{hash_bits}>, '
-    '{n_sparse}, {n_chan}, {pool_size}, {pixel_parallel_factor}, {chan_parallel_factor}>'
+    '{n_sparse}, {n_chan}, {in_height}, {in_width}, {pool_height}, {pool_width}, '
+    '{pixel_parallel_factor}, {chan_parallel_factor}>'
     '({input}, {output}, {hash_in}, {hash_out});'
 )
 
@@ -309,7 +314,10 @@ def format(self, node):
         params = self._default_function_params(node)
         params['n_sparse'] = node.get_attr('n_sparse')
         params['n_chan'] = node.get_attr('n_chan')
-        params['pool_size'] = node.get_attr('pool_size')
+        params['in_height'] = node.get_attr('in_height')
+        params['in_width'] = node.get_attr('in_width')
+        params['pool_height'] = node.get_attr('pool_height')
+        params['pool_width'] = node.get_attr('pool_width')
         params['hash_bits'] = _get_hash_bits(node)
         params['hash_in'] = node.get_attr('hash_in_name')
         params['hash_out'] = node.get_attr('hash_out_name')

hls4ml/converters/keras_v3/sparsepixels.py

@@ -77,6 +77,7 @@ def handle(
         # Clear any stale state from a previous conversion in the same Python process
         _sparse_context.clear()
         _sparse_context['n_sparse'] = n_sparse
+        _sparse_context['spatial'] = (int(in_height), int(in_width))
 
         # Hash stores 1-based H and W coordinates separately (see nnet_sparsepixels.h::sparse_input_reduce).
         # Spatial dims only shrink through the network (pooling), so input H/W bound the required bits.
@@ -175,18 +176,27 @@ def handle(
 
 
 def _sparse_pooling_config(
-    in_tensors: Sequence['KerasTensor'], out_tensors: Sequence['KerasTensor'], pool_size: int, pool_op: str
+    in_tensors: Sequence['KerasTensor'], out_tensors: Sequence['KerasTensor'], pool_size: tuple, pool_op: str
 ) -> dict[str, Any]:
     """Shared config for the average/max sparse pooling handlers (differ only in pool_op)."""
     feat_shape: tuple[int, ...] = in_tensors[0].shape[1:]  # type: ignore
     n_chan = int(feat_shape[-1])
     n_sparse = _sparse_context.get('n_sparse', 0)
+    pool_h, pool_w = int(pool_size[0]), int(pool_size[1])
+
+    # Spatial dims at this layer (tracked from the input reduction through the pooling chain);
+    # the kernel needs them to invalidate pooled cells falling outside the valid output grid.
+    prev_h, prev_w = _sparse_context.get('spatial', (1, 1))
+    _sparse_context['spatial'] = (prev_h // pool_h, prev_w // pool_w)
 
     return {
         'class_name': 'SparsePooling2D',
         'n_sparse': n_sparse,
         'n_chan': n_chan,
-        'pool_size': pool_size,
+        'in_height': prev_h,
+        'in_width': prev_w,
+        'pool_height': pool_h,
+        'pool_width': pool_w,
         'pool_op': pool_op,
     }
 
@@ -200,7 +210,7 @@ def handle(
         in_tensors: Sequence['KerasTensor'],
         out_tensors: Sequence['KerasTensor'],
     ):
-        return _sparse_pooling_config(in_tensors, out_tensors, int(layer.avg_pool.pool_size[0]), 'avg')
+        return _sparse_pooling_config(in_tensors, out_tensors, tuple(layer.avg_pool.pool_size), 'avg')
 
 
 class MaxPooling2DSparseHandler(KerasV3LayerHandler):
@@ -212,4 +222,4 @@ def handle(
         in_tensors: Sequence['KerasTensor'],
         out_tensors: Sequence['KerasTensor'],
     ):
-        return _sparse_pooling_config(in_tensors, out_tensors, int(layer.max_pool.pool_size[0]), 'max')
+        return _sparse_pooling_config(in_tensors, out_tensors, tuple(layer.max_pool.pool_size), 'max')

hls4ml/model/layers.py

@@ -1955,7 +1955,10 @@ class SparsePooling2D(Layer):
     _expected_attributes = [
         Attribute('n_sparse'),
         Attribute('n_chan'),
-        Attribute('pool_size'),
+        Attribute('in_height'),
+        Attribute('in_width'),
+        Attribute('pool_height'),
+        Attribute('pool_width'),
         Attribute('pool_op', value_type=str, default='avg'),  # 'avg' or 'max'
         TypeAttribute('accum'),
     ]

hls4ml/model/optimizer/passes/bit_exact.py

@@ -824,14 +824,15 @@ def _(layer: SparseActivation):
 @_produce_kif.register
 def _(layer: SparsePooling2D):
     k_in, i_in, f_in = get_input_kifs(layer)[0]
-    # Average pooling divides by pool_size^2, which adds ceil(log2(pool_size^2)) fractional bits
-    # (matching standard Pooling2D). Max pooling just selects an input, so the precision is unchanged.
-    pool_size = layer.attributes['pool_size']
+    # Average pooling divides by the pool area, which adds ceil(log2(pool_height * pool_width))
+    # fractional bits (matching standard Pooling2D). Max pooling just selects an input, so the
+    # precision is unchanged.
+    pool_area = layer.attributes['pool_height'] * layer.attributes['pool_width']
     n_chan = layer.attributes['n_chan']
     if layer.attributes.get('pool_op', 'avg') == 'max':
         extra_f = 0
     else:
-        extra_f = int(np.ceil(np.log2(pool_size * pool_size)))
+        extra_f = int(np.ceil(np.log2(pool_area)))
     k_ch = k_in[:n_chan]
     i_ch = i_in[:n_chan]
     f_ch = f_in[:n_chan] + extra_f

hls4ml/model/optimizer/passes/sparsepixels_flatten.py

@@ -34,13 +34,13 @@ def transform(self, model, node):
         cur = inp
         while cur is not None and not isinstance(cur, SparseInputReduce):
             if isinstance(cur, SparsePooling2D):
-                pool_sizes.append(cur.get_attr('pool_size'))
+                pool_sizes.append((cur.get_attr('pool_height'), cur.get_attr('pool_width')))
             cur = cur.get_input_node()
         if cur is None:
             return False
         height, width = cur.get_attr('in_height'), cur.get_attr('in_width')
-        for pool_size in pool_sizes:
-            height, width = height // pool_size, width // pool_size
+        for pool_h, pool_w in pool_sizes:
+            height, width = height // pool_h, width // pool_w
 
         attrs = {
             'n_sparse': n_sparse,

hls4ml/templates/vivado/nnet_utils/nnet_sparsepixels.h

@@ -223,28 +223,38 @@ void sparse_relu(data_T sparse_arr_feat_in[N_sparse * n_chan], res_T sparse_arr_
 // mapping to it (the is_first test); duplicate pixels of the same cell emit 0. The averaging reads
 // only the input array (no scratch mutation), so it is safe to partially unroll. Two independent
 // knobs: pixel_parallel_factor (N_sparse axis) and chan_parallel_factor (n_chan axis).
-template <class data_T, class res_T, class hash_T, class accum_T, int N_sparse, int n_chan, int pool_size,
-          int pixel_parallel_factor = N_sparse, int chan_parallel_factor = n_chan>
+// Pool height and width are independent (asymmetric pooling). A pooled coordinate falling outside
+// the valid output grid (partial window of an odd input dimension under 'valid' pooling) has its
+// features zeroed, matching the dense layer dropping that window; zero-feature pixels are inert in
+// every downstream sparse kernel. The averaging divides by the full pool area as one reciprocal
+// multiply per axis, skipped for a unit axis whose reciprocal 1.0 the fixed-point type cannot hold.
+template <class data_T, class res_T, class hash_T, class accum_T, int N_sparse, int n_chan, int in_height, int in_width,
+          int pool_height, int pool_width, int pixel_parallel_factor = N_sparse, int chan_parallel_factor = n_chan>
 void sparse_pooling_avg(data_T sparse_arr_feat_in[N_sparse * n_chan], res_T sparse_arr_feat_out[N_sparse * n_chan],
                         hash_T sparse_arr_hash_in[N_sparse * 2], hash_T sparse_arr_hash_out[N_sparse * 2]) {
 
-    constexpr double _pool_size_recip_d = 1.0 / double(pool_size);
-    const ap_fixed<10, 0> pool_size_recip = _pool_size_recip_d;
+    constexpr int out_height = in_height / pool_height;
+    constexpr int out_width = in_width / pool_width;
+    // Unsigned reciprocals: 1/2 = 0.5 needs the unsigned [0, 1) range, since a signed
+    // ap_fixed<10, 0> tops out just below 0.5 and would wrap.
+    const ap_ufixed<10, 0> pool_h_recip = 1.0 / double(pool_height); // only used when pool_height > 1
+    const ap_ufixed<10, 0> pool_w_recip = 1.0 / double(pool_width);  // only used when pool_width > 1
 
     int hash_tmp[N_sparse * 2];
     #pragma HLS ARRAY_PARTITION variable = hash_tmp type = complete dim = 0
 ComputePooledLoc:
     for (int i = 0; i < N_sparse; i++) {
         #pragma HLS UNROLL
-        hash_tmp[2 * i] = (sparse_arr_hash_in[2 * i] - 1) / pool_size + 1;
-        hash_tmp[2 * i + 1] = (sparse_arr_hash_in[2 * i + 1] - 1) / pool_size + 1;
+        hash_tmp[2 * i] = (sparse_arr_hash_in[2 * i] - 1) / pool_height + 1;
+        hash_tmp[2 * i + 1] = (sparse_arr_hash_in[2 * i + 1] - 1) / pool_width + 1;
     }
 
 HashOutLoop:
     for (int i_pixel = 0; i_pixel < N_sparse; i_pixel++) {
         #pragma HLS UNROLL factor = pixel_parallel_factor
         int h_out = hash_tmp[2 * i_pixel];
         int w_out = hash_tmp[2 * i_pixel + 1];
+        bool valid = (h_out <= out_height) && (w_out <= out_width);
 
         bool is_first = true;
     FirstCheck:
@@ -270,8 +280,17 @@ void sparse_pooling_avg(data_T sparse_arr_feat_in[N_sparse * n_chan], res_T spar
                     acc += sparse_arr_feat_in[n_chan * j_pixel + i_chan];
                 }
             }
-            sparse_arr_feat_out[n_chan * i_pixel + i_chan] =
-                is_first ? (res_T)(acc * pool_size_recip * pool_size_recip) : (res_T)0;
+            res_T avg;
+            if (pool_height > 1 && pool_width > 1) {
+                avg = (res_T)(acc * pool_h_recip * pool_w_recip);
+            } else if (pool_height > 1) {
+                avg = (res_T)(acc * pool_h_recip);
+            } else if (pool_width > 1) {
+                avg = (res_T)(acc * pool_w_recip);
+            } else {
+                avg = (res_T)acc;
+            }
+            sparse_arr_feat_out[n_chan * i_pixel + i_chan] = (is_first && valid) ? avg : (res_T)0;
         }
         sparse_arr_hash_out[2 * i_pixel] = h_out;
         sparse_arr_hash_out[2 * i_pixel + 1] = w_out;
@@ -282,25 +301,31 @@ void sparse_pooling_avg(data_T sparse_arr_feat_in[N_sparse * n_chan], res_T spar
 // is_first test), but takes the per-channel maximum of the active pixels in the cell, floored at 0
 // to match dense max pooling over the zero-masked window. Two independent knobs:
 // pixel_parallel_factor (N_sparse axis) and chan_parallel_factor (n_chan axis).
-template <class data_T, class res_T, class hash_T, int N_sparse, int n_chan, int pool_size,
-          int pixel_parallel_factor = N_sparse, int chan_parallel_factor = n_chan>
+// Pool height and width are independent (asymmetric pooling); out-of-range pooled coordinates are
+// zeroed as in the average version.
+template <class data_T, class res_T, class hash_T, int N_sparse, int n_chan, int in_height, int in_width, int pool_height,
+          int pool_width, int pixel_parallel_factor = N_sparse, int chan_parallel_factor = n_chan>
 void sparse_pooling_max(data_T sparse_arr_feat_in[N_sparse * n_chan], res_T sparse_arr_feat_out[N_sparse * n_chan],
                         hash_T sparse_arr_hash_in[N_sparse * 2], hash_T sparse_arr_hash_out[N_sparse * 2]) {
 
+    constexpr int out_height = in_height / pool_height;
+    constexpr int out_width = in_width / pool_width;
+
     int hash_tmp[N_sparse * 2];
     #pragma HLS ARRAY_PARTITION variable = hash_tmp type = complete dim = 0
 ComputePooledLoc:
     for (int i = 0; i < N_sparse; i++) {
         #pragma HLS UNROLL
-        hash_tmp[2 * i] = (sparse_arr_hash_in[2 * i] - 1) / pool_size + 1;
-        hash_tmp[2 * i + 1] = (sparse_arr_hash_in[2 * i + 1] - 1) / pool_size + 1;
+        hash_tmp[2 * i] = (sparse_arr_hash_in[2 * i] - 1) / pool_height + 1;
+        hash_tmp[2 * i + 1] = (sparse_arr_hash_in[2 * i + 1] - 1) / pool_width + 1;
     }
 
 HashOutLoop:
     for (int i_pixel = 0; i_pixel < N_sparse; i_pixel++) {
         #pragma HLS UNROLL factor = pixel_parallel_factor
         int h_out = hash_tmp[2 * i_pixel];
         int w_out = hash_tmp[2 * i_pixel + 1];
+        bool valid = (h_out <= out_height) && (w_out <= out_width);
 
         bool is_first = true;
     FirstCheck:
@@ -327,7 +352,7 @@ void sparse_pooling_max(data_T sparse_arr_feat_in[N_sparse * n_chan], res_T spar
                     vmax = v;
                 }
             }
-            sparse_arr_feat_out[n_chan * i_pixel + i_chan] = is_first ? (res_T)vmax : (res_T)0;
+            sparse_arr_feat_out[n_chan * i_pixel + i_chan] = (is_first && valid) ? (res_T)vmax : (res_T)0;
         }
         sparse_arr_hash_out[2 * i_pixel] = h_out;
         sparse_arr_hash_out[2 * i_pixel + 1] = w_out;

test/pytest/test_sparsepixels.py

@@ -95,6 +95,42 @@ def test_sparse_cnn(test_case_id, backend, pool):
     _convert_and_check(model, x, test_root_path / test_case_id, backend)
 
 
+def _build_sparse_cnn_chain(input_shape, n, pools, pool='avg', threshold=0.4):
+    """Conv+pool blocks with per-block (asymmetric) pool sizes, e.g. x-only pooling on odd heights."""
+    iq_conf = QuantizerConfig(place='datalane', q_type='kif', i0=4, f0=8, overflow_mode='WRAP')
+    pool_cls = MaxPooling2DSparse if pool == 'max' else AveragePooling2DSparse
+    with (
+        QuantizerConfigScope(place='all', default_q_type='kbi', overflow_mode='SAT_SYM'),
+        QuantizerConfigScope(place='datalane', default_q_type='kif', overflow_mode='WRAP'),
+        LayerConfigScope(enable_ebops=True, enable_iq=True, beta0=1e-5),
+    ):
+        x_in = keras.Input(shape=input_shape, name='x_in')
+        x, keep_mask = InputReduce(n=n, threshold=threshold, name='input_reduce')(x_in)
+        for k, ps in enumerate(pools, 1):
+            x = QConv2DSparse(
+                filters=2, kernel_size=3, name=f'conv{k}', padding='same', strides=1, activation='relu', iq_conf=iq_conf
+            )([x, keep_mask])
+            x, keep_mask = pool_cls(ps, name=f'pool{k}')([x, keep_mask])
+        x = Flatten(name='flatten')(x)
+        x = QDense(1, name='dense', iq_conf=iq_conf)(x)
+    return keras.Model(x_in, x, name='cnn_sparse_asym_test')
+
+
+@pytest.mark.parametrize('backend', ['Vivado', 'Vitis'])
+@pytest.mark.parametrize('pool', ['avg', 'max'])
+@pytest.mark.parametrize('pools', [((1, 4), (2, 2)), ((2, 2), (1, 2))])
+def test_sparse_cnn_asymmetric_pool(test_case_id, backend, pool, pools):
+    # Asymmetric pool sizes on an odd input height (9 rows). The ((2, 2), ...) chain pools the odd
+    # dimension first, so pixels of the dropped partial window (out-of-range pooled coordinates)
+    # flow through a later conv and pool: they must be invalidated to match the dense keras layers.
+    np.random.seed(44)
+    keras.utils.set_random_seed(44)
+
+    model = _build_sparse_cnn_chain(input_shape=(9, 16, 1), n=6, pools=pools, pool=pool)
+    x = _make_sparse_inputs(n_samples=1000, h=9, w=16, n_active_per_sample=5)
+    _convert_and_check(model, x, test_root_path / test_case_id, backend)
+
+
 @pytest.mark.parametrize('backend', ['Vitis'])
 def test_sparse_cnn_parallelization(test_case_id, backend):
     # Partial parallelization and the streaming input reduce only change the unroll/implementation,