DistributeGeneric: Wrong secretness analysis on multiple loop carried variable

[ZenithalHourlyRate]

Followup of #1606

For the input

module {
  func.func @hv_matmul(%arg0: !secret.secret<tensor<1x1024xf32>>) -> !secret.secret<tensor<1x1024xf32>> {
    %c1 = arith.constant 1 : index
    %0 = secret.generic ins(%arg0 : !secret.secret<tensor<1x1024xf32>>) attrs = {__argattrs = [{mgmt.mgmt = #mgmt.mgmt<level = 1>}], __resattrs = [{mgmt.mgmt = #mgmt.mgmt<level = 0>}]} {
    ^body(%input0: tensor<1x1024xf32>):
      %1:2 = affine.for %arg1 = 1 to 1024 iter_args(%arg2 = %input0, %arg3 = %input0) -> (tensor<1x1024xf32>, tensor<1x1024xf32>) {
        %3 = tensor_ext.rotate %arg3, %c1 {mgmt.mgmt = #mgmt.mgmt<level = 1>} : tensor<1x1024xf32>, index
        %4 = arith.addf %arg2, %3 {mgmt.mgmt = #mgmt.mgmt<level = 1>} : tensor<1x1024xf32>
        affine.yield %4, %3 : tensor<1x1024xf32>, tensor<1x1024xf32>
      } {__argattrs = [{mgmt.mgmt = #mgmt.mgmt<level = 1>}, {mgmt.mgmt = #mgmt.mgmt<level = 1>}], mgmt.mgmt = #mgmt.mgmt<level = 1>}
      %2 = mgmt.modreduce %1#0 {mgmt.mgmt = #mgmt.mgmt<level = 0>} : tensor<1x1024xf32>
      secret.yield %2 : tensor<1x1024xf32>
    } -> !secret.secret<tensor<1x1024xf32>>
    return %0 : !secret.secret<tensor<1x1024xf32>>
  }
}

After applying --secret-distribute-generic, in the intermediate state, it will become a valid IR like

"func.func"() <{arg_attrs = [{mgmt.mgmt = #mgmt.mgmt<level = 1>}], function_type = (!secret.secret<tensor<1x1024xf32>>) -> !secret.secret<tensor<1x1024xf32>>, res_attrs = [{mgmt.mgmt = #mgmt.mgmt<level = 0>}], sym_name = "hv_matmul"}> ({
^bb0(%arg0: !secret.secret<tensor<1x1024xf32>>):
  %0 = "arith.constant"() <{value = 1 : index}> : () -> index
  %1:2 = "affine.for"(%arg0, %arg0) <{lowerBoundMap = affine_map<() -> (1)>, operandSegmentSizes = array<i32: 0, 0, 2>, step = 1 : index, upperBoundMap = affine_map<() -> (1024)>}> ({
  ^bb0(%arg2: index, %arg3: !secret.secret<tensor<1x1024xf32>>, %arg4: !secret.secret<tensor<1x1024xf32>>):
    %4 = "secret.generic"() ({
      %7 = "tensor_ext.rotate"(%arg4, %0) {mgmt.mgmt = #mgmt.mgmt<level = 1>} : (!secret.secret<tensor<1x1024xf32>>, index) -> tensor<1x1024xf32>
      "secret.yield"(%7) : (tensor<1x1024xf32>) -> ()
    }) : () -> !secret.secret<tensor<1x1024xf32>>
    %5 = "secret.generic"(%arg3, %4) ({
    ^bb0(%arg5: tensor<1x1024xf32>, %arg6: tensor<1x1024xf32>):
      %6 = "arith.addf"(%arg5, %arg6) <{fastmath = #arith.fastmath<none>}> {mgmt.mgmt = #mgmt.mgmt<level = 1>} : (tensor<1x1024xf32>, tensor<1x1024xf32>) -> tensor<1x1024xf32>
      "secret.yield"(%6) : (tensor<1x1024xf32>) -> ()
    }) : (!secret.secret<tensor<1x1024xf32>>, !secret.secret<tensor<1x1024xf32>>) -> !secret.secret<tensor<1x1024xf32>>
    "affine.yield"(%5, %4) : (!secret.secret<tensor<1x1024xf32>>, !secret.secret<tensor<1x1024xf32>>) -> ()
  }) {__argattrs = [{mgmt.mgmt = #mgmt.mgmt<level = 1>}, {mgmt.mgmt = #mgmt.mgmt<level = 1>}], mgmt.mgmt = #mgmt.mgmt<level = 1>} : (!secret.secret<tensor<1x1024xf32>>, !secret.secret<tensor<1x1024xf32>>) -> (!secret.secret<tensor<1x1024xf32>>, !secret.secret<tensor<1x1024xf32>>)
  %2 = "secret.generic"(%1#0) ({
  ^bb0(%arg1: tensor<1x1024xf32>):
    %3 = "mgmt.modreduce"(%arg1) {mgmt.mgmt = #mgmt.mgmt<level = 0>} : (tensor<1x1024xf32>) -> tensor<1x1024xf32>
    "secret.yield"(%3) : (tensor<1x1024xf32>) -> ()
  }) {__resattrs = [{mgmt.mgmt = #mgmt.mgmt<level = 0>}]} : (!secret.secret<tensor<1x1024xf32>>) -> !secret.secret<tensor<1x1024xf32>>
  "func.return"(%2) : (!secret.secret<tensor<1x1024xf32>>) -> ()
}) : () -> ()

Then, canonicalization pattern (CollapseSecretlessGeneric) will think the tensor_ext.rotate's secret.generic is a secretless generic, and it will be Collapsed into

  %1:2 = "affine.for"(%arg0, %arg0) <{lowerBoundMap = affine_map<() -> (1)>, operandSegmentSizes = array<i32: 0, 0, 2>, step = 1 : index, upperBoundMap = affine_map<() -> (1024)>}> ({
  ^bb0(%arg2: index, %arg3: !secret.secret<tensor<1x1024xf32>>, %arg4: !secret.secret<tensor<1x1024xf32>>):
    %4 = "tensor_ext.rotate"(%arg4, %0) {mgmt.mgmt = #mgmt.mgmt<level = 1>} : (!secret.secret<tensor<1x1024xf32>>, index) -> tensor<1x1024xf32>
    %5 = "secret.generic"(%arg3, %4) ({
    ^bb0(%arg5: tensor<1x1024xf32>, %arg6: tensor<1x1024xf32>):
      %6 = "arith.addf"(%arg5, %arg6) <{fastmath = #arith.fastmath<none>}> {mgmt.mgmt = #mgmt.mgmt<level = 1>} : (tensor<1x1024xf32>, tensor<1x1024xf32>) -> tensor<1x1024xf32>
      "secret.yield"(%6) : (tensor<1x1024xf32>) -> ()
    }) : (!secret.secret<tensor<1x1024xf32>>, tensor<1x1024xf32>) -> !secret.secret<tensor<1x1024xf32>>
    "affine.yield"(%5, %4) : (!secret.secret<tensor<1x1024xf32>>, tensor<1x1024xf32>) -> ()

which is clearly incorrect.

If we replace tensor_ext.rotate with other sentences like %3 = arith.addf %arg3, %arg3 {mgmt.mgmt = #mgmt.mgmt<level = 1>} : tensor<1x1024xf32>, the same error happens. Note that %arg3 is the second loop carried variable.

If we replace that sentence into %3 = arith.addf %arg2, %arg2 {mgmt.mgmt = #mgmt.mgmt<level = 1>} : tensor<1x1024xf32>, which is the first loop carried variable, then secret analysis is happy.