diff --git a/hls4ml/backends/vivado/passes/sparsepixels.py b/hls4ml/backends/vivado/passes/sparsepixels.py index 6b8a065e22..ef7c5bf648 100644 --- a/hls4ml/backends/vivado/passes/sparsepixels.py +++ b/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') diff --git a/hls4ml/converters/keras_v3/sparsepixels.py b/hls4ml/converters/keras_v3/sparsepixels.py index 781db02164..9ce03abf90 100644 --- a/hls4ml/converters/keras_v3/sparsepixels.py +++ b/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') diff --git a/hls4ml/model/layers.py b/hls4ml/model/layers.py index 08945876f6..d571196656 100644 --- a/hls4ml/model/layers.py +++ b/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'), ] diff --git a/hls4ml/model/optimizer/passes/bit_exact.py b/hls4ml/model/optimizer/passes/bit_exact.py index 18a5abc0e1..e073bd6fa2 100644 --- a/hls4ml/model/optimizer/passes/bit_exact.py +++ b/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 diff --git a/hls4ml/model/optimizer/passes/sparsepixels_flatten.py b/hls4ml/model/optimizer/passes/sparsepixels_flatten.py index f07a26927a..9cd56ef58c 100644 --- a/hls4ml/model/optimizer/passes/sparsepixels_flatten.py +++ b/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, diff --git a/hls4ml/templates/vivado/nnet_utils/nnet_sparsepixels.h b/hls4ml/templates/vivado/nnet_utils/nnet_sparsepixels.h index c31a516d52..a3e2034781 100644 --- a/hls4ml/templates/vivado/nnet_utils/nnet_sparsepixels.h +++ b/hls4ml/templates/vivado/nnet_utils/nnet_sparsepixels.h @@ -223,21 +223,30 @@ 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 +// 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 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: @@ -245,6 +254,7 @@ void sparse_pooling_avg(data_T sparse_arr_feat_in[N_sparse * n_chan], res_T spar #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,18 +301,23 @@ 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 +// Pool height and width are independent (asymmetric pooling); out-of-range pooled coordinates are +// zeroed as in the average version. +template 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: @@ -301,6 +325,7 @@ void sparse_pooling_max(data_T sparse_arr_feat_in[N_sparse * n_chan], res_T spar #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; diff --git a/test/pytest/test_sparsepixels.py b/test/pytest/test_sparsepixels.py index 53f30c4e5c..e4ab9bb45d 100644 --- a/test/pytest/test_sparsepixels.py +++ b/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,