ASM: suggest OpResult name for BGV/CKKS/Openfhe

[ZenithalHourlyRate]

This utilizes the OpAsmOpInterface's getAsmResultNames， just similar to how getAlias works in #1131

  func.func @dot_product(%arg0: !openfhe.crypto_context, %arg1: !ct_L2_, %arg2: !ct_L2_) -> !ct_L0_ {
    %cst = arith.constant dense<[0, 0, 0, 0, 0, 0, 0, 1]> : tensor<8xi64>
    %ct = openfhe.mul_no_relin %arg0, %arg1, %arg2 : (!openfhe.crypto_context, !ct_L2_, !ct_L2_) -> !ct_L2_D3_
    %ct_0 = openfhe.relin %arg0, %ct : (!openfhe.crypto_context, !ct_L2_D3_) -> !ct_L2_
    %ct_1 = openfhe.rot %arg0, %ct_0 {index = 4 : index} : (!openfhe.crypto_context, !ct_L2_) -> !ct_L2_
    %ct_2 = openfhe.add %arg0, %ct_0, %ct_1 : (!openfhe.crypto_context, !ct_L2_, !ct_L2_) -> !ct_L2_
    %ct_3 = openfhe.rot %arg0, %ct_2 {index = 2 : index} : (!openfhe.crypto_context, !ct_L2_) -> !ct_L2_
    %ct_4 = openfhe.add %arg0, %ct_2, %ct_3 : (!openfhe.crypto_context, !ct_L2_, !ct_L2_) -> !ct_L2_
    %ct_5 = openfhe.rot %arg0, %ct_4 {index = 1 : index} : (!openfhe.crypto_context, !ct_L2_) -> !ct_L2_
    %ct_6 = openfhe.add %arg0, %ct_4, %ct_5 : (!openfhe.crypto_context, !ct_L2_, !ct_L2_) -> !ct_L2_
    %ct_7 = openfhe.mod_reduce %arg0, %ct_6 : (!openfhe.crypto_context, !ct_L2_) -> !ct_L1_
    %pt = openfhe.make_packed_plaintext %arg0, %cst : (!openfhe.crypto_context, tensor<8xi64>) -> !pt
    %ct_8 = openfhe.mul_plain %arg0, %ct_7, %pt : (!openfhe.crypto_context, !ct_L1_, !pt) -> !ct_L1_
    %ct_9 = openfhe.rot %arg0, %ct_8 {index = 7 : index} : (!openfhe.crypto_context, !ct_L1_) -> !ct_L1_
    %ct_10 = lwe.reinterpret_underlying_type %ct_9 : !ct_L1_ to !ct_L1_1
    %ct_11 = openfhe.mod_reduce %arg0, %ct_10 : (!openfhe.crypto_context, !ct_L1_1) -> !ct_L0_
    return %ct_11 : !ct_L0_
  }
  func.func @dot_product__encrypt__arg0(%arg0: !openfhe.crypto_context, %arg1: tensor<8xi16>, %arg2: !openfhe.public_key) -> !ct_L2_ {
    %0 = arith.extsi %arg1 : tensor<8xi16> to tensor<8xi64>
    %pt = openfhe.make_packed_plaintext %arg0, %0 : (!openfhe.crypto_context, tensor<8xi64>) -> !pt
    %ct = openfhe.encrypt %arg0, %pt, %arg2 : (!openfhe.crypto_context, !pt, !openfhe.public_key) -> !ct_L2_
    return %ct : !ct_L2_
  }
  func.func @dot_product__encrypt__arg1(%arg0: !openfhe.crypto_context, %arg1: tensor<8xi16>, %arg2: !openfhe.public_key) -> !ct_L2_ {
    %0 = arith.extsi %arg1 : tensor<8xi16> to tensor<8xi64>
    %pt = openfhe.make_packed_plaintext %arg0, %0 : (!openfhe.crypto_context, tensor<8xi64>) -> !pt
    %ct = openfhe.encrypt %arg0, %pt, %arg2 : (!openfhe.crypto_context, !pt, !openfhe.public_key) -> !ct_L2_
    return %ct : !ct_L2_
  }
  func.func @dot_product__decrypt__result0(%arg0: !openfhe.crypto_context, %arg1: !ct_L0_, %arg2: !openfhe.private_key) -> i16 {
    %pt = openfhe.decrypt %arg0, %arg1, %arg2 : (!openfhe.crypto_context, !ct_L0_, !openfhe.private_key) -> !pt1
    %0 = lwe.rlwe_decode %pt {encoding = #full_crt_packing_encoding, ring = #ring_Z65537_i64_1_x8_} : !pt1 -> i16
    return %0 : i16
  }
  func.func @dot_product__generate_crypto_context() -> !openfhe.crypto_context {
    %params = openfhe.gen_params  {insecure = false, mulDepth = 2 : i64, plainMod = 4295294977 : i64} : () -> !openfhe.cc_params
    %cc = openfhe.gen_context %params {supportFHE = false} : (!openfhe.cc_params) -> !openfhe.crypto_context
    return %cc : !openfhe.crypto_context
  }
  func.func @dot_product__configure_crypto_context(%arg0: !openfhe.crypto_context, %arg1: !openfhe.private_key) -> !openfhe.crypto_context {
    openfhe.gen_mulkey %arg0, %arg1 : (!openfhe.crypto_context, !openfhe.private_key) -> ()
    openfhe.gen_rotkey %arg0, %arg1 {indices = array<i64: 1, 2, 4, 7>} : (!openfhe.crypto_context, !openfhe.private_key) -> ()
    return %arg0 : !openfhe.crypto_context
  }

[ZenithalHourlyRate]

Unfortunately, I was not able to alter the function argument name, as func::FuncOp is not in our control so its interface implementation is fixed.

I tried the following hack

  struct FuncOpHeirInterface
      : public ::mlir::OpAsmOpInterface::ExternalModel<FuncOpHeirInterface,
                                                       func::FuncOp> {
    void getAsmBlockArgumentNames(Operation *op, Region &region,
                                  ::mlir::OpAsmSetValueNameFn setNameFn) const {
      for (auto &block : region) {
        for (auto arg : block.getArguments()) {
          setNameFn(arg, "test");
        }
      }
    }
  };

  registry.addExtension(+[](MLIRContext *ctx, func::FuncDialect *dialect) {
    func::FuncOp::attachInterface<FuncOpHeirInterface>(*ctx);
  });

But it did not work as we will get Ignoring repeated interface registration, because func::FuncOp already registered the interface and the ExternalModel registration should not work.

This is intended behavior, as https://mlir.llvm.org/docs/Interfaces/ has suggested

Note: It is strongly encouraged to only use this mechanism if you “own” the interface being externally applied. This prevents a situation where neither the owner of the dialect containing the object nor the owner of the interface are aware of an interface implementation, which can lead to duplicate or diverging implementations.

[AlexanderViand-Intel]

Thanks, this looks super nice!

Unfortunately, I was not able to alter the function argument name, as func::FuncOp is not in our control so its interface implementation is fixed.

I wonder if there's appetite upstream for changingfunc.func's printer so that it respects some (new? existing?) name-defining interface on Types (if present for a given type) and uses that, rather than arg? However, I could forsee this breaking a lot of tests, both upstream and in various downstream projects, since lots of them hardcoded %argN names in.

[ZenithalHourlyRate]

I get a hack for it. Just override the default implementation at the linker level.

// hack here: another template specialization for FuncOp
// expect linker to pick this one
template <>
void ::mlir::detail::OpAsmOpInterfaceInterfaceTraits::
    Model<mlir::func::FuncOp>::getAsmBlockArgumentNames(
        mlir::detail::OpAsmOpInterfaceInterfaceTraits::Concept const *,
        mlir::Operation *op, mlir::Region &region,
        ::mlir::OpAsmSetValueNameFn setNameFn) {
  ::mlir::heir::getAsmBlockArgumentNames(op, region, setNameFn);
}

Refactored the code so much cleaner now. Should refactor our uses of OpAsmOpInterface and OpAsmDialectInterface so we can reuse the suggestNameForType (even reuse it in SelectVariableNames).

Now we have much informative IR

func.func @dot_product(%cc: !openfhe.crypto_context, %ct: !ct_L2_, %ct_0: !ct_L2_) -> !ct_L0_ 
...
func.func @dot_product__configure_crypto_context(%cc: !openfhe.crypto_context, %skey: !openfhe.private_key) -> !openfhe.crypto_context {
    openfhe.gen_mulkey %cc, %skey : (!openfhe.crypto_context, !openfhe.private_key) -> ()
    openfhe.gen_rotkey %cc, %skey {indices = array<i64: 1, 2, 4, 7>} : (!openfhe.crypto_context, !openfhe.private_key) -> ()
    return %cc : !openfhe.crypto_context
  }

[AlexanderViand-Intel]

Refactored the code so much cleaner now. Should refactor our uses of OpAsmOpInterface and OpAsmDialectInterface so we can reuse the suggestNameForType (even reuse it in SelectVariableNames).

Nice, that part LGTM!

[AlexanderViand-Intel]

I'm really torn on this one - I love the IR it produces, but this is pretty hacky. At the very least, I'd say lets expand the comment here, maybe link it back to this PR or a new issue that explains why the other alternatives don't work.

[j2kun]

I'm OK with it, but I would be quick to remove it if it causes problems.

As long as no tests depend on the naming convention, it should be purely cosmetic.

[AlexanderViand-Intel]

LGTM with the disclaimer comment, but I'd love to hear @j2kun and/or @asraa also weigh in on the "hack" part.

(Which might not happen before the new year, since people are out for the holidays)

[ZenithalHourlyRate]

I wonder if there's appetite upstream for changingfunc.func's printer so that it respects some (new? existing?) name-defining interface on Types (if present for a given type) and uses that, rather than arg? However, I could forsee this breaking a lot of tests, both upstream and in various downstream projects, since lots of them hardcoded %argN names in.

I have thought of ways to do this in upstream so we can avoid using the hack (should be a long-term upstreaming process though).

Now the Op/Dialect themselves can control the implementation of the interface, but other people can not (similar to OpAsmDialectInterface?).

My first thinking is to allow user to override the Interface implementation (detailed as OpAsmOpInterfaceInterfaceTraits::Model/FallbackModel/ExternalModel and InterfaceMap), but the order of overriding becomes a issue (who takes the precedence, what if there are two overriding and things become undeterministic, depending on how people order code).

Another more reasonable way is that the default implementation of OpAsmOpInterface provides API for CallbackFn (similar to TypeConverter) so users can get more fine grained control as often not all Values one Op uses are from its own Dialect. This way is much more extensible but needs careful designing.

Or, since in our implementation it depends on the result type, maybe a default OpAsmInterface implementation could consult the Type with something interface like TypeAsmInterface and can get name from it.

I wonder how the upstream would think about it, maybe I should open a thread on LLVM discourse.

[j2kun]

Suggested change

	Op<BGV_Dialect, mnemonic, traits # [OpAsmOpInterface]> {
	Op<BGV_Dialect, mnemonic, traits # [DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>]> {

[j2kun]

This will define the methods for you, so you can remove the extra class declaration below. Cf. https://github.com/llvm/llvm-project/blob/1557eeda738d7dbe51d2f52fce28a1fd6f5844ce/mlir/include/mlir/Dialect/Mesh/IR/MeshOps.td#L85 for an upstream example

Same for the other td files.

[ZenithalHourlyRate]

I put the getAsmResultNames definition here because otherwise I have to instantiate it for each op in Ops.cpp

[j2kun]

Suggested change

	.Case<NewLWEPublicKeyType>([&](Type) { return "pkey"; })
	.Case<NewLWESecretKeyType>([&](Type) { return "skey"; })
	.Case<NewLWEPublicKeyType>([&](Type) { return "pk"; })
	.Case<NewLWESecretKeyType>([&](Type) { return "sk"; })

[j2kun]

Same below

[j2kun]

I'm OK with it, but I would be quick to remove it if it causes problems.

As long as no tests depend on the naming convention, it should be purely cosmetic.