Skip to content

Implementing secrecy aware operation rebalancing#2917

Open
aymane-eljerari wants to merge 1 commit intogoogle:mainfrom
aymane-eljerari:mul-add-rebalance
Open

Implementing secrecy aware operation rebalancing#2917
aymane-eljerari wants to merge 1 commit intogoogle:mainfrom
aymane-eljerari:mul-add-rebalance

Conversation

@aymane-eljerari
Copy link
Copy Markdown

@aymane-eljerari aymane-eljerari commented May 3, 2026

Add support for secrecy aware operation rebalancing using the SecretnessAnalysis solver. Addressing the following issue: (#836)

Before rebuilding the compute graph in tryBalanceBlock() the code partitions the operands based on their retreived Secretness Value (public values are placed at the beginning of the list maitaining the original DFS order after partitioning).

I updated the balance_sigmoid.mlir since it was failing due to the re-arragement of operands. The new transformed output maintains sigmoid's computational correctness.

I added a new secrecy_aware_balance.mlir tests. I takes a plaintext function argument and plaintext arith.constants.

@google-cla
Copy link
Copy Markdown

google-cla Bot commented May 3, 2026

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

Copy link
Copy Markdown
Collaborator

@asraa asraa left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

looks great!

// CHECK-DAG: %[[CONST_0:.*]] = arith.constant dense<5.000000e-01> : tensor<1x16xf32>
// CHECK-DAG: %[[CONST_1:.*]] = arith.constant dense<5.000000e-02> : tensor<1x16xf32>

// CHECK: %[[RET:.*]] = secret.generic(%arg0: !secret.secret<tensor<1x16xf32>>, %arg1: !secret.secret<tensor<1x16xf32>>, %arg2: tensor<1x16xf32>)
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
// CHECK: %[[RET:.*]] = secret.generic(%arg0: !secret.secret<tensor<1x16xf32>>, %arg1: !secret.secret<tensor<1x16xf32>>, %arg2: tensor<1x16xf32>)
// CHECK: %[[RET:.*]] = secret.generic(%{{[^:]*}}: !secret.secret<tensor<1x16xf32>>, %{{[^:]*}}: !secret.secret<tensor<1x16xf32>>, %{{[^:]*}}: tensor<1x16xf32>)

@j2kun
Copy link
Copy Markdown
Collaborator

j2kun commented May 5, 2026

When you're ready please squash the commits into one and I will start the automation to merge it

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants