Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Annotate ModuleOp scheme/backend attribute #1586

Merged
merged 1 commit into from
Mar 21, 2025

Conversation

ZenithalHourlyRate
Copy link
Collaborator

Follows #1420 (comment). The goal is making an "unmanaged" pipeline for openfhe so the loop support can work without all the mgmt nuances.

To do that I used the least managed strategy for B/FV and reuse it for BGV/CKKS when encountered. To get the multiplicativeDepth property for openfhe, the MulResultAnalysis becomes MulDepthAnalysis, which is different from level analysis as the former relies on mulop and the later relies on modreduce.

Parameter generation will just generate the fallback parameter when seeing the openfhe flag.

@ZenithalHourlyRate
Copy link
Collaborator Author

Though for the current loop support we will run into assertion...

bazel run //tools:heir-opt -- $PWD/tests/Examples/common/halevi_shoup_matmul.mlir --annotate-module="backend=openfhe" --mlir-to-ckks="ciphertext-degree=16 experimental-disable-loop-unroll=true"

The IR can survive all the way through linalg-to-tensor-ext and mgmt passes, and get

    %0 = secret.generic ins(%arg0 : !secret.secret<tensor<1x16xf32>>) attrs = {__argattrs = [{mgmt.mgmt = #mgmt.mgmt<level = 1>}], __resattrs = [{mgmt.mgmt = #mgmt.mgmt<level = 1>}]} {
    ^body(%input0: tensor<1x16xf32>):
      %1 = arith.mulf %input0, %cst {mgmt.mgmt = #mgmt.mgmt<level = 1>} : tensor<1x16xf32>
      %2 = arith.addf %1, %cst_1 {mgmt.mgmt = #mgmt.mgmt<level = 1>} : tensor<1x16xf32>
      %3:2 = affine.for %arg1 = 1 to 16 iter_args(%arg2 = %2, %arg3 = %input0) -> (tensor<1x16xf32>, tensor<1x16xf32>) {
        %4 = tensor_ext.rotate %arg3, %c1 {mgmt.mgmt = #mgmt.mgmt<level = 1>} : tensor<1x16xf32>, index
        %extracted_slice = tensor.extract_slice %cst_0[%arg1, 0] [1, 16] [1, 1] : tensor<16x16xf32> to tensor<1x16xf32>
        %5 = arith.mulf %4, %extracted_slice {mgmt.mgmt = #mgmt.mgmt<level = 1>} : tensor<1x16xf32>
        %6 = arith.addf %arg2, %5 {mgmt.mgmt = #mgmt.mgmt<level = 1>} : tensor<1x16xf32>
        affine.yield %6, %4 : tensor<1x16xf32>, tensor<1x16xf32>
      } {__argattrs = [{mgmt.mgmt = #mgmt.mgmt<level = 1>}, {mgmt.mgmt = #mgmt.mgmt<level = 1>}], mgmt.mgmt = #mgmt.mgmt<level = 1>}
      secret.yield %3#0 : tensor<1x16xf32>
    } -> !secret.secret<tensor<1x16xf32>>

Which is certainly something Openfhe is happy with. Then DistributeGeneric is unhappy somewhere

heir-opt: lib/Dialect/Secret/Transforms/DistributeGeneric.cpp:345
Assertion `isa<BlockArgument>(value) && "not sure what to do here, file a bug"' failed.

I think this should be addressed in another issue and should not affect this PR. The loop produced by this PR is valid.

@j2kun j2kun added the pull_ready Indicates whether a PR is ready to pull. The copybara worker will import for internal testing label Mar 21, 2025
@j2kun
Copy link
Collaborator

j2kun commented Mar 21, 2025

I think there may be some places where the top level getOperation() in the pass is not pointing to a module op (e.g., if you run a pass pipeline that anchors on func.func) but this can be fixed when we run into it.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
pull_ready Indicates whether a PR is ready to pull. The copybara worker will import for internal testing
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants