Add fold-plaintext-masks for simple cleanup after implement-shift-network

[j2kun]

Fixes #2337

The sole idea here is to take repeated applications of plaintext one-hot masks and combine them via bitwise-and'ing the entries of the masks.

The example test still produces some inefficiencies:

$ python ~/fhe/heir/scripts/lit_to_bazel.py --run=True tests/Regressionissue_2337.mlir
module {
  func.func @trivial_insert(%arg0: !secret.secret<tensor<2x32xi32>>) -> !secret.secret<tensor<2x32xi32>> {
    %cst = arith.constant dense<[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]> : tensor<32xi32>
    %cst_0 = arith.constant dense<0> : tensor<32xi32>
    %cst_1 = arith.constant dense<[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]> : tensor<32xi32>
    %c16 = arith.constant 16 : index
    %0 = tensor.empty() : tensor<2x32xi32>
    %1 = secret.generic(%arg0: !secret.secret<tensor<2x32xi32>>) {
    ^body(%input0: tensor<2x32xi32>):
      %extracted_slice = tensor.extract_slice %input0[1, 0] [1, 32] [1, 1] : tensor<2x32xi32> to tensor<32xi32>
      %2 = arith.muli %extracted_slice, %cst : tensor<32xi32>
      %3 = tensor_ext.rotate %2, %c16 : tensor<32xi32>, index
      %4 = arith.muli %3, %cst_1 : tensor<32xi32>
      %5 = arith.addi %4, %2 : tensor<32xi32>
      %6 = arith.muli %extracted_slice, %cst_1 : tensor<32xi32>
      %7 = tensor_ext.rotate %6, %c16 : tensor<32xi32>, index
      %8 = arith.addi %6, %7 : tensor<32xi32>
      %9 = arith.addi %5, %8 : tensor<32xi32>
      %inserted_slice = tensor.insert_slice %9 into %0[0, 0] [1, 32] [1, 1] : tensor<32xi32> into tensor<2x32xi32>
      %inserted_slice_2 = tensor.insert_slice %cst_0 into %inserted_slice[1, 0] [1, 32] [1, 1] : tensor<32xi32> into tensor<2x32xi32>
      secret.yield %inserted_slice_2 : tensor<2x32xi32>
    } -> !secret.secret<tensor<2x32xi32>>
    return %1 : !secret.secret<tensor<2x32xi32>>
  }
}

These three lines show the problem

%2 = arith.muli %extracted_slice, %cst : tensor<32xi32>
%3 = tensor_ext.rotate %2, %c16 : tensor<32xi32>, index
%4 = arith.muli %3, %cst_1 : tensor<32xi32>

It extracts the first half, rotates it to the second half, then masks the second half (a no-op) before adding it to a copy of the first half. That second mask can be eliminated because the other half of the entries can be inferred to be zero. Something for a future slot analysis to handle...

[asraa]

thank you!