heir icon indicating copy to clipboard operation
heir copied to clipboard

tf2fhe: Figure out the correct support for `memref` globals in the pipeline

Open asraa opened this issue 2 years ago • 1 comments

Right now, the tf2fhe pipeline from #267 is blocked on the following issue from hello_world, see this paste https://gist.github.com/asraa/97e95df7b42b84bcf738e6ae5434da9d for the partial debug output.

    %0:6 = secret.generic {
      %14 = memref.get_global @__constant_16x1xi8 : memref<16x1xi8>
      %15 = memref.get_global @__constant_16xi32 : memref<16xi32>
      %16 = memref.get_global @__constant_16x16xi8 : memref<16x16xi8>
      %17 = memref.get_global @__constant_16xi32_0 : memref<16xi32>
      %18 = memref.get_global @__constant_1x16xi8 : memref<1x16xi8>
      %alloc = memref.alloc() {alignment = 64 : i64} : memref<1x16xi8>
      secret.yield %14, %15, %16, %17, %18, %alloc : memref<16x1xi8>, memref<16xi32>, memref<16x16xi8>, memref<16xi32>, memref<1x16xi8>, memref<1x16xi8>
    } -> (!secret.secret<memref<16x1xi8>>, !secret.secret<memref<16xi32>>, !secret.secret<memref<16x16xi8>>, !secret.secret<memref<16xi32>>, !secret.secret<memref<1x16xi8>>, !secret.secret<memref<1x16xi8>>)
    affine.for %arg1 = 0 to 1 {
      affine.for %arg2 = 0 to 16 {
        secret.generic ins(%0#0, %0#5 : !secret.secret<memref<16x1xi8>>, !secret.secret<memref<1x16xi8>>) {
        ^bb0(%arg3: memref<16x1xi8>, %arg4: memref<1x16xi8>):
          %14 = affine.load %arg3[%arg2, %arg1] : memref<16x1xi8>
          affine.store %14, %arg4[%arg1, %arg2] : memref<1x16xi8>
          secret.yield
        }
      }
    }

In this, you can see memref.get_globals that have been wrapped in a secret.generic.

These are plain contstants, and so may not need this wrapping, so perhaps some data flow analysis is needed for distribute-secret-generic, or I can also assume that these values are to be trivially encrypted and hoisted outside the generic.

The current choke point is that we "assume" we must yosys-optimize this no-input region...

asraa avatar Dec 14 '23 17:12 asraa

Ideas:

  • Only run yosys-optimizer when the body has some arithmetic instruction that we need to booleanize - seems fragile? Or only when there are translate bodies (memref ops are translateable, but some of these have multiple return vals)?
  • Ops like memref.get_global, affine.load and affine.store don't really need to be wrapped in a secret.generic body. Should they just be hoisted out of the body? I think we would only trivially encrypt the loaded values that are involved in the ciphertext ops (the logic from comb-to-cggi does this)
  • I think the former relies on not wrapping the get_globals in secret values on yield, which seems to require the dataflow analysis...
  • Canonicalize to add the globals into the scope so that translation won't fail, but this seems bad too because translating back to MLIR would be messy

The main goal for me seems to be to get the trivial ops out of the secret.generic bodies. I think the right way to do that is to do a dataflow analysis (since trivial can also mean plaintext-plaintext ops, which we currently don't have but may in the future)

asraa avatar Dec 14 '23 17:12 asraa