AnnotateMgmt for programs with loops fails - LevelAnalysis yields infinite loop

[asraa]

Root cause: Analyses do not understand region bearing ops (for loops) properly, we need an implementation of visitNonControlFlowArguments most likely.

Reproducer:

$ heir-opt --secret-insert-mgmt-ckks
  func.func @loop(%arg0: !secret.secret<tensor<1x1024xf32>>) -> !secret.secret<tensor<1x1024xf32>> {
    %0 = secret.generic ins(%arg0 : !secret.secret<tensor<1x1024xf32>>) {
    ^body(%input0: tensor<1x1024xf32>):
      %1 = affine.for %arg1 = 0 to 3 iter_args(%arg3 = %input0) -> (tensor<1x1024xf32>) {
        %3 = arith.mulf %arg3, %arg3 : tensor<1x1024xf32>
        affine.yield %3 : tensor<1x1024xf32>
      }
      secret.yield %1 : tensor<1x1024xf32>
    } -> !secret.secret<tensor<1x1024xf32>>
    return %0 : !secret.secret<tensor<1x1024xf32>>
  }

This program has a single ctxt-ctxt multiplication inside, and the result should have a relinearization inserted after the mul so that the block arg and yielded value always have dimension 2. The pass also inserts modreduce before the multiplication even when first mul option is off because it seems that the args are from a mulresult too.

So I think LevelAnalysis, DimensionAnalysis, and MulResultAnalysis all need to be aware of the region bearing operation.

When you run the reproducer, it will run into an infinite loop when it re-runs the solver after the relinearization insertion. I added some debugs to the level analysis and saw it continuining to increment the level of multiplication over and over in the loop.

LevelAnalysis: Visiting operation: %2 = mgmt.modreduce %arg2 : tensor<1x1024xf32>
propogate: %2 = mgmt.modreduce %arg2 : tensor<1x1024xf32>is changed 1 to state 2769
LevelAnalysis: Visiting operation: %3 = arith.mulf %2, %2 : tensor<1x1024xf32>
propogate: %3 = arith.mulf %2, %2 : tensor<1x1024xf32>is changed 1 to state 2769
LevelAnalysis: Visiting operation: %3 = arith.mulf %2, %2 : tensor<1x1024xf32>
propogate: %3 = arith.mulf %2, %2 : tensor<1x1024xf32>is changed 0 to state 2769
LevelAnalysis: Visiting operation: %4 = mgmt.relinearize %3 : tensor<1x1024xf32>
propogate: %4 = mgmt.relinearize %3 : tensor<1x1024xf32>is changed 1 to state 2769
LevelAnalysis: visitNonControlFlowArguments: %1 = affine.for %arg1 = 0 to 3 iter_args(%arg2 = %input0) -> (tensor<1x1024xf32>) {
  %2 = mgmt.modreduce %arg2 : tensor<1x1024xf32>
  %3 = arith.mulf %2, %2 : tensor<1x1024xf32>
  %4 = mgmt.relinearize %3 : tensor<1x1024xf32>
  affine.yield %4 : tensor<1x1024xf32>
}
LevelAnalysis: visitNonControlFlowArguments: %1 = affine.for %arg1 = 0 to 3 iter_args(%arg2 = %input0) -> (tensor<1x1024xf32>) {
  %2 = mgmt.modreduce %arg2 : tensor<1x1024xf32>
  %3 = arith.mulf %2, %2 : tensor<1x1024xf32>
  %4 = mgmt.relinearize %3 : tensor<1x1024xf32>
  affine.yield %4 : tensor<1x1024xf32>
}
LevelAnalysis: Visiting operation: %2 = mgmt.modreduce %arg2 : tensor<1x1024xf32>
propogate: %2 = mgmt.modreduce %arg2 : tensor<1x1024xf32>is changed 1 to state 2770
LevelAnalysis: Visiting operation: %3 = arith.mulf %2, %2 : tensor<1x1024xf32>
propogate: %3 = arith.mulf %2, %2 : tensor<1x1024xf32>is changed 1 to state 2770
LevelAnalysis: Visiting operation: %3 = arith.mulf %2, %2 : tensor<1x1024xf32>
propogate: %3 = arith.mulf %2, %2 : tensor<1x1024xf32>is changed 0 to state 2770
LevelAnalysis: Visiting operation: %4 = mgmt.relinearize %3 : tensor<1x1024xf32>
propogate: %4 = mgmt.relinearize %3 : tensor<1x1024xf32>is changed 1 to state 2770
LevelAnalysis: visitNonControlFlowArguments: %1 = affine.for %arg1 = 0 to 3 iter_args(%arg2 = %input0) -> (tensor<1x1024xf32>) {
  %2 = mgmt.modreduce %arg2 : tensor<1x1024xf32>
  %3 = arith.mulf %2, %2 : tensor<1x1024xf32>
  %4 = mgmt.relinearize %3 : tensor<1x1024xf32>
  affine.yield %4 : tensor<1x1024xf32>
}
LevelAnalysis: visitNonControlFlowArguments: %1 = affine.for %arg1 = 0 to 3 iter_args(%arg2 = %input0) -> (tensor<1x1024xf32>) {
  %2 = mgmt.modreduce %arg2 : tensor<1x1024xf32>
  %3 = arith.mulf %2, %2 : tensor<1x1024xf32>
  %4 = mgmt.relinearize %3 : tensor<1x1024xf32>
  affine.yield %4 : tensor<1x1024xf32>
}
LevelAnalysis: Visiting operation: %2 = mgmt.modreduce %arg2 : tensor<1x1024xf32>
propogate: %2 = mgmt.modreduce %arg2 : tensor<1x1024xf32>is changed 1 to state 2771
LevelAnalysis: Visiting operation: %3 = arith.mulf %2, %2 : tensor<1x1024xf32>
propogate: %3 = arith.mulf %2, %2 : tensor<1x1024xf32>is changed 1 to state 2771
LevelAnalysis: Visiting operation: %3 = arith.mulf %2, %2 : tensor<1x1024xf32>

cc @ZenithalHourlyRate

[j2kun]

continuining to increment the level of multiplication over and over in the loop

I think this is what happens naturally when the lattice doesn't have a "top." For instance, if you do something like an "integer constant value analysis" where the lattice tracks a subset of possible values an integer variable could be (e.g., {3, 5, 7} + {0, 1} -> {3, 4, 5, 6, 7} because you could add the zero or one to each value from the first set), then the set will just keep growing bigger and bigger over a loop. So to compensate you'd declare that any set larger than 10 elements has as its lattice element a sentinel value called "top" which represents formally "all possible values" and informally, "I give up". Then anything joined with top becomes top and it's an absorbing steady state that allows the worklist to reach a fixed point.

I'm guessing these analyses just keep incrementing through the loop, and for whatever reason the sparse analysis doesn't understand that a loop with static bounds would only have a finite number of iterations.

I'm starting to think more and more that we should fork/rewrite our own dataflow analysis framework that has special cases for FHE, for example when loops have static iteration bounds. But in the mean time this could be fixed by generalizing the lattice to have a "top" value where any level larger than, say, 25 becomes "top" and similar for others whose lattice elements are unbound integers.

[ZenithalHourlyRate]

See also #1181, which is also an infinite loop when running the analysis. But the situation is different.

The pass also inserts modreduce before the multiplication even when first mul option is off because it seems that the args are from a mulresult too.

The analysis would find the loop carried variable as a previous mult result so it will insert modreduce.

The modreduce insertion is correct in that if we do full loop unroll, we have this result, which is expected.

      %1 = arith.mulf %input0, %input0 {mgmt.mgmt = #mgmt.mgmt<level = 3, dimension = 3>} : tensor<1x1024xf32>
      %2 = mgmt.relinearize %1 {mgmt.mgmt = #mgmt.mgmt<level = 3>} : tensor<1x1024xf32>
      %3 = mgmt.modreduce %2 {mgmt.mgmt = #mgmt.mgmt<level = 2>} : tensor<1x1024xf32>
      %4 = arith.mulf %3, %3 {mgmt.mgmt = #mgmt.mgmt<level = 2, dimension = 3>} : tensor<1x1024xf32>
      %5 = mgmt.relinearize %4 {mgmt.mgmt = #mgmt.mgmt<level = 2>} : tensor<1x1024xf32>
      %6 = mgmt.modreduce %5 {mgmt.mgmt = #mgmt.mgmt<level = 1>} : tensor<1x1024xf32>
      %7 = arith.mulf %6, %6 {mgmt.mgmt = #mgmt.mgmt<level = 1, dimension = 3>} : tensor<1x1024xf32>
      %8 = mgmt.relinearize %7 {mgmt.mgmt = #mgmt.mgmt<level = 1>} : tensor<1x1024xf32>
      %9 = mgmt.modreduce %8 {mgmt.mgmt = #mgmt.mgmt<level = 0>} : tensor<1x1024xf32>
      secret.yield %9 : tensor<1x1024xf32>

Notice that, before the first arith.mulf there is no modreduce, so the correct way to express the loop is

%1 = affine.for %arg1 = 0 to 3 iter_args(%arg2 = %input0) -> (tensor<1x1024xf32>) {
  %3 = arith.mulf %arg2, %arg2 : tensor<1x1024xf32>
  %4 = mgmt.relinearize %3 : tensor<1x1024xf32>
  %5 = mgmt.modreduce %4 : tensor<1x1024xf32>
  affine.yield %5 : tensor<1x1024xf32>
}

instead of

%1 = affine.for %arg1 = 0 to 3 iter_args(%arg2 = %input0) -> (tensor<1x1024xf32>) {
  %2 = mgmt.modreduce %arg2 : tensor<1x1024xf32>
  %3 = arith.mulf %2, %2 : tensor<1x1024xf32>
  %4 = mgmt.relinearize %3 : tensor<1x1024xf32>
  affine.yield %4 : tensor<1x1024xf32>
}

But even if we can express it in arith dialect, when lowered further to LWE dialect, such loop is not expressable, as we have %arg2 of type !ct_L3, !ct_L2, !ct_L1, !ct_L0 as the loop progresses.

%input0 = ... : !ct_L3
%1 = affine.for %arg1 = 0 to 3 iter_args(%arg2 = %input0) -> (!ct_L0) {
  %3 = arith.mulf %arg2, %arg2 : !what type
  %4 = mgmt.relinearize %3 : !what type
  %5 = mgmt.modreduce %4 : !what type
  affine.yield %5 : !what type
}

The conclusion is that such invariant breaking loop is not expressible in FHE world, and the analysis should reject it as it could not be lowered.

For matmul_loop.mlir, we could keep an invariant (level the same, dimension always 2 so LWE type the same) if we insert mgmt in the following way

    %1 = affine.for %arg2 = 0 to 1 iter_args(%arg3 = %arg1) -> (tensor<1x16xf32>) {
      %2 = affine.for %arg4 = 0 to 16 iter_args(%arg5 = %arg3) -> (tensor<1x16xf32>) {
        %3 = affine.for %arg6 = 0 to 16 iter_args(%arg7 = %arg5) -> (tensor<1x16xf32>) {
          %extracted = tensor.extract %arg0[%arg2, %arg6] : tensor<1x16xf32>
          %extracted_0 = tensor.extract %0[%arg6, %arg4] : tensor<16x16xf32>
          %extracted_1 = tensor.extract %arg7[%arg2, %arg4] : tensor<1x16xf32>
          // ct-pt mul, dimension always 2
          %4 = arith.mulf %extracted, %extracted_0 : f32
          %5 = arith.addf %extracted_1, %4 : f32
          // no modreduce to keep the level invariant
          %inserted = tensor.insert %5 into %arg7[%arg2, %arg4] : tensor<1x16xf32>
          affine.yield %inserted : tensor<1x16xf32>
        }
        affine.yield %3 : tensor<1x16xf32>
      }
      affine.yield %2 : tensor<1x16xf32>
    }

So that the loop carried variable has stable LWE type.

I agree that we "need an implementation of visitNonControlFlowArguments", but we first need to analysis the loop body to check whether a loop invariant could be established here to make sure such loop is expressible in FHE world.