While lowering an AlexNet-Tiny-style CKKS model, linalg-fuse-linear-ops incorrectly rewrites a square activation after linalg.matvec.
The pass treats the secret-dependent matvec output as a public scale and fuses it into the matvec weights, producing weights that depend on the secret input.
Reproducer
%l1 = linalg.matvec ins(%w1, %x : tensor<256x768xf32>, tensor<768xf32>)
outs(%b1 : tensor<256xf32>) -> tensor<256xf32>
%sq1 = arith.mulf %l1, %l1 : tensor<256xf32>
Run:
bazel run //tools:heir-opt --
--linalg-canonicalizations --linalg-fuse-linear-ops
alexnet-tiny.mlir
Actual behavior
The pass rewrites the square into something equivalent to:
%l1 = linalg.matvec ins(%w1, %x) outs(%b1) -> tensor<256xf32>
%broadcasted = linalg.broadcast ins(%l1) ...
%scaled_w1 = arith.mulf %w1, %broadcasted
%new = linalg.matvec ins(%scaled_w1, %x) outs(%b1) -> tensor<256xf32>
This is invalid because %scaled_w1 now depends on the secret input.
Expected behavior
The square activation should remain an elementwise ciphertext-ciphertext multiplication. It should not be folded into the linear layer weights.
Root cause
findLinearOpAndOperand in LinalgFuseLinearOps.cpp checks whether either operand of arith.mulf is produced by a linalg::LinalgOp, but it does not check whether the other operand is public / constant / non-secret.
For arith.mulf %l1, %l1, both operands are the same secret-dependent matvec result, so the pass incorrectly treats one side as the linear op result and the other side as a scale.
Suggested fix
Reject the pattern if the proposed scale operand is the same SSA value as the linear op result, or more generally if it is secret-dependent / produced by a linalg op.
While lowering an AlexNet-Tiny-style CKKS model,
linalg-fuse-linear-opsincorrectly rewrites a square activation afterlinalg.matvec.The pass treats the secret-dependent matvec output as a public scale and fuses it into the matvec weights, producing weights that depend on the secret input.
Reproducer
%l1 = linalg.matvec ins(%w1, %x : tensor<256x768xf32>, tensor<768xf32>)
outs(%b1 : tensor<256xf32>) -> tensor<256xf32>
%sq1 = arith.mulf %l1, %l1 : tensor<256xf32>
Run:
bazel run //tools:heir-opt --
--linalg-canonicalizations --linalg-fuse-linear-ops
alexnet-tiny.mlir
Actual behavior
The pass rewrites the square into something equivalent to:
%l1 = linalg.matvec ins(%w1, %x) outs(%b1) -> tensor<256xf32>
%broadcasted = linalg.broadcast ins(%l1) ...
%scaled_w1 = arith.mulf %w1, %broadcasted
%new = linalg.matvec ins(%scaled_w1, %x) outs(%b1) -> tensor<256xf32>
This is invalid because %scaled_w1 now depends on the secret input.
Expected behavior
The square activation should remain an elementwise ciphertext-ciphertext multiplication. It should not be folded into the linear layer weights.
Root cause
findLinearOpAndOperand in LinalgFuseLinearOps.cpp checks whether either operand of arith.mulf is produced by a linalg::LinalgOp, but it does not check whether the other operand is public / constant / non-secret.
For arith.mulf %l1, %l1, both operands are the same secret-dependent matvec result, so the pass incorrectly treats one side as the linear op result and the other side as a scale.
Suggested fix
Reject the pattern if the proposed scale operand is the same SSA value as the linear op result, or more generally if it is secret-dependent / produced by a linalg op.