[Draft] Scheme selection passes

[kragall]

This is a draft of how scheme selection could be handled with HEIR on a very abstract level. The idea is that this could be helpful to non-expert users and could be a stepping stone towards dealing with scheme switching, see the discussion in #2019.

[asraa]

I agree with Alex that this would need to run at a high enough level where we haven't yet run approximation / arithmetization.
At those high levels I expect we'll also see linalg operations / matmul operations that would be arithmetic operations (and heavily weighted to arithmetic schemes)

[asraa]

in a scheme selection, i would think that if arith ops have smaller bit width (2-8 bits), the cost of evaluation in CGGI like schemes will be lower than high bit width arithmetic (8+ bits). i think this info doesn't synthesize this but maybe the cost model can take this into consideration?

[AlexanderViand]

I think that's a great point, since arithmetic width should be pretty essential to any reasonable cost model!

[AlexanderViand]

I think the core question here is really what the cost model needs to look like. I think the current bool/bit/arith ops counter approach is too simplistic to be useful, yet at the same time I don't think we need some super complex model based on lots of benchmarks/etc.

I think the hardest part is to find some way to track roughly what kind of mult-depth/bootstrapping obligation a program induces (when viewed as an arithmetic problem). Accounting for the cost of basic arithmetic (considering bitwidth, as mentioned by Asra) should be relatively straightfoward, especially if we make some simplifying assumptions on roughly how expensive a matmul/etc will be (which does depend on the kernel, but I don't think the cost model needs to be all that accurate).

Since really, the main question for HEIR right now is "do we do binarization or do we do arithmetization", we can probably just track two accountings of the costs, one assuming we went for CGGI and one assuming we went for CKKS/BGV/BFV and ignore the finer points inside those approaches until HEIR actually supports, e.g, circuit bootstrapped CGGI or interpolation-based airhtmetization for BFV/BGV.

Btw, I don't think its very important for the cost model to be "correct" at the moment, it just needs to be reasonable, i.e., don't pick CGGI for a program with 12 64-bit matmuls just because it sees one comparison at the end, or pick CKKS with N=64k and 1000-bit modulus for a program that only operates on like 4 small inputs in total, just because it's all add/mul ops.

[AlexanderViand]

@kragall I see the "PR format requirements" failed - you should follow the steps at https://heir.dev/docs/development/ and run pre-commit run --all-files to clean those up, and probably take a look at https://heir.dev/docs/development/ide/ to make sure your IDE formats things correctly automatically.

[kragall]

Sorry, forgot about pre-commit when I switched to my new laptop. Now it's reinstalled and the checks should succeed

[kragall]

The current SchemeInfoAnalysis annotates the func operation with the different types of ops that occur inside the function. The SchemeSelectionAnalysis is currently empty, because I can't get it to visit the func op.

You can run both analysis by using --annotate-scheme-info, e.g. on tests/Examples/common/dot_product_8.mlir . If you use the --debug --debug-only=SchemeSelection flags, the SchemeSelectionAnalysis prints out all ops that it visits.

[j2kun]

The current SchemeInfoAnalysis annotates the func operation with the different types of ops that occur inside the function. The SchemeSelectionAnalysis is currently empty, because I can't get it to visit the func op.

You can run both analysis by using --annotate-scheme-info, e.g. on tests/Examples/common/dot_product_8.mlir . If you use the --debug --debug-only=SchemeSelection flags, the SchemeSelectionAnalysis prints out all ops that it visits.

There must be some unpushed changes because heir-opt doesn't build on your branch. It's missing a BUILD file in lib/Analysis/SchemeInfoAnalysis

[j2kun]

Adding the build file and otherwise it seems to work for me:

$ bazel run //tools:heir-opt -- --annotate-scheme-info --debug --debug-only=SchemeInfo $PWD/tests/Examples/common/dot_product_8.mlir
Visiting: arith.constant. Visiting: arith.constant. Visiting: tensor.extract. Visiting: tensor.extract. Visiting: arith.muli. Counting: NatureOfComputation(numBoolOps=0; numBitOps=0; numIntArithOps=1; numRealArithOps=0; numCmpOps=0; numNonLinOps=0)
Visiting: arith.addi. Counting: NatureOfComputation(numBoolOps=0; numBitOps=0; numIntArithOps=1; numRealArithOps=0; numCmpOps=0; numNonLinOps=0)
Visiting: arith.addi. Counting: NatureOfComputation(numBoolOps=0; numBitOps=0; numIntArithOps=1; numRealArithOps=0; numCmpOps=0; numNonLinOps=0)
Counting here: arith.constant
Counting here: arith.constant
Counting here: affine.for
Counting here: tensor.extract
Counting here: tensor.extract
Counting here: arith.muli
Counting here: arith.addi
Counting here: affine.yield
Counting here: func.return
Writing annotations here: func.func
module {
  func.func @dot_product(%arg0: tensor<8xi16> {secret.secret}, %arg1: tensor<8xi16> {secret.secret}) -> i16 attributes {natcomp.bitOps = 0 : i64, natcomp.boolOps = 0 : i64, natcomp.cmpOps = 0 : i64, natcomp.intOps = 3 : i64, natcomp.nonLinOps = 0 : i64, natcomp.realOps = 0 : i64} {
    %c0 = arith.constant 0 : index
    %c0_i16 = arith.constant 0 : i16
    %0 = affine.for %arg2 = 0 to 8 iter_args(%arg3 = %c0_i16) -> (i16) {
      %extracted = tensor.extract %arg0[%arg2] : tensor<8xi16>
      %extracted_0 = tensor.extract %arg1[%arg2] : tensor<8xi16>
      %1 = arith.muli %extracted, %extracted_0 {natcomp.intOps = 1 : i64} : i16
      %2 = arith.addi %arg3, %1 {natcomp.intOps = 1 : i64} : i16
      affine.yield %2 : i16
    }
    return %0 : i16
  }
}

[j2kun]

Perhaps one thing to note is that the analysis won't "visit" your func.func directly because the data flow framework handles that by another means (initializing all the func args to the entry state to start the worklist)

[kragall]

So visitOperation is not supposed to visit func.func. Does this mean that the best course of action would be to add a function in SchemeInfoAnalysis that's called after the annotate function and add a walk through the module?

[j2kun]

I'm not sure I understand the problem. The function is annotated in the example above. What outcome are you trying to achieve?

[kragall]

After annotating func.func similar to the example above, the next step in my opinion would be to use the annotation for actually selecting a scheme. I wanted to do this in the SchemeSelectionAnalysis and use either a call to the annotate-module pass or a similar annotation as the result of this analysis. But for the analysis to work I need to access the annotations of func.func and I tried to do this similar to the other analyses I've come across, i.e. by using a type-switch statement in visitOperation to gather the information about all func.func ops in the lattice and then run the analysis in a separate function.

[j2kun]

Ah yes, you won't be able to do it in an analysis using the data flow framework. But you could just do the analysis work in the pass after the analysis is run, as you said with a walk. You can also just create any sort of class object you want and call it from inside runOnOperation.