[Feature]: Standalone convolution operation

[chart21]

Feature Request Type

Usability

Have you searched existing issues?

Yes

Is your feature request related to a problem?

Since the paper proposes significant improvements in linear layer computation, I would be interested to isolate the convolution operation to benchmark it against related work.

Describe features you want to add to SPU

Could you give me a pointer to the relevant code that I would need to specify convolutions parameters (input dimension, number of filters, kernel dimensions, stride, padding etc.) and batch size to run a standalone 2PC convolution with the desired parameters?

For instance, is this the function I need?

spu::Value Convolution2D(SPUContext *ctx, const spu::Value &input,
                         const spu::Value &kernel,
                         const ConvolutionConfig &config,
                         const Shape &result_shape);

And I assume there is some setup code I require to run a full distributed 2PC with that function.

Describe features you want to add to SPU

[fionser]

I would prefer to run the code in Python rather than using the C++ interface.
Once you built the whole project, you can use any JAX API to simulate the 2PC computation.
For example (add a new python to examples/python/microbench/conv2d.py)

import unittest

import jax.nn as jnn
import jax.numpy as jnp
import numpy as np
from jax import lax

import spu.intrinsic as si
import spu.spu_pb2 as spu_pb2
import spu.utils.simulation as ppsim

# Perform convolution


def conv2d():
    config = spu_pb2.RuntimeConfig(
        protocol=spu_pb2.ProtocolKind.CHEETAH, field=spu_pb2.FieldType.FM64
    )
    config.enable_hal_profile = True
    config.experimental_enable_colocated_optimization = False
    config.cheetah_2pc_config.enable_mul_lsb_error = True
    config.cheetah_2pc_config.approx_less_precision = 4
    config.max_concurrency = 2

    sim = ppsim.Simulator(2, config)

    # Define a simple edge-detection kernel
    kernel = jnp.array([[1, 1, 0], [1, 0, -1], [0, -1, -1]], dtype=jnp.float32)[
        :, :, None, None
    ]
    # Create a synthetic image
    image = np.zeros((1, 200, 200, 1), dtype=jnp.float32)
    image[0, 50:150, 50:150, 0] = 1.0

    def _con2d(img, ker):
        output = lax.conv_general_dilated(
            img,
            ker,
            window_strides=(1, 1),
            padding="SAME",
            dimension_numbers=("NHWC", "HWIO", "NHWC"),
        )
        return output

    spu_fn = ppsim.sim_jax(sim, _con2d)
    z = spu_fn(image, kernel)
    g = _con2d(image, kernel)
    diff = z - g

    print("conv2d max diff = {}".format(np.max(diff)))


if __name__ == "__main__":
    conv2d()

[chart21]

Thanks a lot, the example works for me. How do I get the best performance?
Compared to OpenCheetah I get worse communication and runtime. But I would expect much better results given Bumblebee's improvements.

I tried running the first convolution of ResNet50 with 32-bit numbers as shown below:

   # ResNet-50 first convolution 
    kernel = jnp.empty((7, 7, 3, 64), dtype=jnp.float32)

    # ImageNet-sized image
    image = np.zeros((1, 224, 224, 3), dtype=jnp.float32)
    image[0, 50:150, 50:150, :] = 1.0

    def _con2d(img, ker):
        output = lax.conv_general_dilated(
            img,
            ker,
            window_strides=(2, 2),
            padding="SAME",
            dimension_numbers=("NHWC", "HWIO", "NHWC"),
        )
        return output

and used the following settings

   config.enable_hal_profile = True
    config.experimental_enable_colocated_optimization = True
    config.cheetah_2pc_config.enable_mul_lsb_error = True
    config.cheetah_2pc_config.approx_less_precision = 4
    config.max_concurrency = 32

and got the following results. I would expect around an order of magnitude less communication and runtime.

[2025-07-28 13:30:13.806] [info] [thread_pool.cc:30] Create a fixed thread pool with size 39
[2025-07-28 13:30:14.263] [info] [cheetah_dot.cc:310] CheetahDot uses 2@1 modulus 4096 degree for 32 bit ring (packing=disabled)
[2025-07-28 13:30:28.905] [info] [cheetah_dot.cc:475] 1@12544x147x64 => 1x64x64 Recv 0.166 MiB, Response 393.092 MiB Pack 0 ms (none)
[2025-07-28 13:30:29.645] [info] [cheetah_dot.cc:475] 1@12544x147x64 => 1x64x64 Recv 0.166 MiB, Response 393.093 MiB Pack 0 ms (none)
[2025-07-28 13:30:43.014] [info] [api.cc:165] [Profiling] SPU execution _con2d completed, input processing took 2.682e-06s, execution took 29.148602883s, output processing took 1.2713e-05s, total time 29.148618278s.
[2025-07-28 13:30:43.029] [info] [api.cc:215] HAL profiling: total time 28.764767619
[2025-07-28 13:30:43.029] [info] [api.cc:218] - f_tensordot, executed 1 times, duration 28.764702774s, send MiB 416.5534 recv MiB 407.7575
[2025-07-28 13:30:43.029] [info] [api.cc:218] - seal, executed 1 times, duration 6.4845e-05s, send MiB 0 recv MiB 0
[2025-07-28 13:30:43.029] [info] [api.cc:215] MPC profiling: total time 29.070256818
[2025-07-28 13:30:43.029] [info] [api.cc:218] - mmul_aa, executed 1 times, duration 15.903810975s, send MiB 393.3139 recv MiB 393.3136
[2025-07-28 13:30:43.029] [info] [api.cc:218] - trunc_a, executed 1 times, duration 12.857116705s, send MiB 23.2394 recv MiB 14.4439
[2025-07-28 13:30:43.029] [info] [api.cc:218] - concatenate, executed 1 times, duration 0.137507651s, send MiB 0 recv MiB 0
[2025-07-28 13:30:43.029] [info] [api.cc:218] - extract_slice, executed 12544 times, duration 0.123232267s, send MiB 0 recv MiB 0
[2025-07-28 13:30:43.029] [info] [api.cc:218] - reshape, executed 12549 times, duration 0.039257688s, send MiB 0 recv MiB 0
[2025-07-28 13:30:43.029] [info] [api.cc:218] - pad, executed 1 times, duration 0.008672595s, send MiB 0 recv MiB 0
[2025-07-28 13:30:43.029] [info] [api.cc:218] - transpose, executed 2 times, duration 0.000602243s, send MiB 0 recv MiB 0
[2025-07-28 13:30:43.029] [info] [api.cc:218] - p2a, executed 1 times, duration 5.6694e-05s, send MiB 0 recv MiB 0
[2025-07-28 13:30:43.029] [info] [api.cc:228] Link details: total send MiB 416.5534, recv MiB 407.7575, send actions 83300

[fionser]

However, the BumbleBee paper does not optimize the Conv2D operation. Actually, we ommit the optimized implementation in OpenCheetah due to the difficulty to re-implement the Cheetah's Conv2D in the SPU lib.
Your pick of ResNet-50 the first convolution with 7 kernel size is worse for the current BumbleBee's code :(

Also the current python script will secret share both input tensor and weight matrix and resulting an 2x costs.
We can see that mmul_aa the API is called (aka matrix-mult-arith_share-arith_share) . For the OpenCheetah, the weight matrix is not secretly shared.

To setup the Client/Server run in OpenBumbleBee, you can checkout the examples

    spu_input_ids = ppd.device("P1")(lambda x: x)(input_ids) # Party-1 provides the input tensor
    spu_params = ppd.device("P2")(lambda x: x)(params) # Party-2 provides the weight matrix

Then the config experimental_enable_colocated_optimization: true will kick-in, and we optimze the mmul_aa to mmul_av (aka matrix-multiplication-arith_share-private)