heir-translate won't ignore non-entry-point functions and errors since they have not been booleanized

[kevaundray]

Aim

I am trying to go from a c program to tfhe-bool-rust.

Issue

An unexpected error occurs in heir-translate, which I believe is due to the non-entrypoint function not being booleanized or removed. This is a guess from the error that gets outputted. Below I give instructions on how to reproduce this.

The C Program

// main.cc
int foo(int x) {
  return x * x;
}

int isqrt(int num) {
  int res = 0;
  int bit = 1 << 14;  // ((unsigned) INT16_MAX + 1) / 2.

  for (int i = 0; i < 8; ++i) {
    if (num >= res + bit) {
      num -= res + bit;
      res = (res >> 1) + bit;
    } else {
      num >>= 1;
    }
    bit >>= 2;
  }
  return foo(res);
}

This was copied from Jeremy's polygeist PR and modified to have the foo function call.

Note: The one in Jeremy's PR required me to unroll the loops for it eventually be translated to tfhe-rust-bool.

Polygeist

I convert it from C to standard MLIR using polygeist with the following call:

cgeist -function=* -raise-scf-to-affine -S -O3 main.cc > main.mlir

Polygeist output

module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<i32, dense<32> : vector<2xi32>>, #dlti.dl_entry<i8, dense<8> : vector<2xi32>>, #dlti.dl_entry<i16, dense<16> : vector<2xi32>>, #dlti.dl_entry<!llvm.ptr, dense<64> : vector<4xi32>>, #dlti.dl_entry<i1, dense<8> : vector<2xi32>>, #dlti.dl_entry<f80, dense<128> : vector<2xi32>>, #dlti.dl_entry<!llvm.ptr<272>, dense<64> : vector<4xi32>>, #dlti.dl_entry<i64, dense<64> : vector<2xi32>>, #dlti.dl_entry<!llvm.ptr<271>, dense<32> : vector<4xi32>>, #dlti.dl_entry<f64, dense<64> : vector<2xi32>>, #dlti.dl_entry<f16, dense<16> : vector<2xi32>>, #dlti.dl_entry<!llvm.ptr<270>, dense<32> : vector<4xi32>>, #dlti.dl_entry<f128, dense<128> : vector<2xi32>>, #dlti.dl_entry<"dlti.endianness", "little">, #dlti.dl_entry<"dlti.stack_alignment", 128 : i32>>, llvm.data_layout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128", llvm.target_triple = "x86_64-unknown-linux-gnu", "polygeist.target-cpu" = "x86-64", "polygeist.target-features" = "+cmov,+cx8,+fxsr,+mmx,+sse,+sse2,+x87", "polygeist.tune-cpu" = "generic"} {
  func.func @_Z3fooi(%arg0: i32) -> i32 attributes {llvm.linkage = #llvm.linkage<external>} {
    %0 = arith.muli %arg0, %arg0 : i32
    return %0 : i32
  }
  func.func @_Z5isqrti(%arg0: i32) -> i32 attributes {llvm.linkage = #llvm.linkage<external>} {
    %c2_i32 = arith.constant 2 : i32
    %c1_i32 = arith.constant 1 : i32
    %c16384_i32 = arith.constant 16384 : i32
    %c0_i32 = arith.constant 0 : i32
    %0:3 = affine.for %arg1 = 0 to 8 iter_args(%arg2 = %c16384_i32, %arg3 = %c0_i32, %arg4 = %arg0) -> (i32, i32, i32) {
      %2 = arith.addi %arg3, %arg2 : i32
      %3 = arith.cmpi sge, %arg4, %2 : i32
      %4:2 = scf.if %3 -> (i32, i32) {
        %6 = arith.subi %arg4, %2 : i32
        %7 = arith.shrsi %arg3, %c1_i32 : i32
        %8 = arith.addi %7, %arg2 : i32
        scf.yield %8, %6 : i32, i32
      } else {
        %6 = arith.shrsi %arg4, %c1_i32 : i32
        scf.yield %arg3, %6 : i32, i32
      }
      %5 = arith.shrsi %arg2, %c2_i32 : i32
      affine.yield %5, %4#0, %4#1 : i32, i32, i32
    }
    %1 = arith.muli %0#1, %0#1 : i32
    return %1 : i32
  }
}

The attributes seem to cause heir-translate to fail, so I remove them.

The cleaned up Polygeist output

// main.mlir
module {
  func.func @_Z3fooi(%arg0: i32) -> i32 {
    %0 = arith.muli %arg0, %arg0 : i32
    return %0 : i32
  }
  func.func @_Z5isqrti(%arg0: i32) -> i32 {
    %c2_i32 = arith.constant 2 : i32
    %c1_i32 = arith.constant 1 : i32
    %c16384_i32 = arith.constant 16384 : i32
    %c0_i32 = arith.constant 0 : i32
    %0:3 = affine.for %arg1 = 0 to 8 iter_args(%arg2 = %c16384_i32, %arg3 = %c0_i32, %arg4 = %arg0) -> (i32, i32, i32) {
      %2 = arith.addi %arg3, %arg2 : i32
      %3 = arith.cmpi sge, %arg4, %2 : i32
      %4:2 = scf.if %3 -> (i32, i32) {
        %6 = arith.subi %arg4, %2 : i32
        %7 = arith.shrsi %arg3, %c1_i32 : i32
        %8 = arith.addi %7, %arg2 : i32
        scf.yield %8, %6 : i32, i32
      } else {
        %6 = arith.shrsi %arg4, %c1_i32 : i32
        scf.yield %arg3, %6 : i32, i32
      }
      %5 = arith.shrsi %arg2, %c2_i32 : i32
      affine.yield %5, %4#0, %4#1 : i32, i32, i32
    }
    %1 = arith.muli %0#1, %0#1 : i32
    return %1 : i32
  }
}

The above is what I feed into heir-opt.

Heir-opt

I run the following bash script that repeatedly calls heir-opt:

heir-opt --inline main.mlir > main-inline.mlir
heir-opt --remove-dead-values main-inline.mlir >  main-inline-dce.mlir
heir-opt --affine-loop-unroll="unroll-full" main-inline-dce.mlir >  main-unroll.mlir
heir-opt --secretize=entry-function=_Z5isqrti --wrap-generic main-unroll.mlir >  main-secret-unroll.mlir
heir-opt --convert-if-to-select main-secret-unroll.mlir >  main-secret-unroll-if-select.mlir

heir-opt --yosys-optimizer='mode=Boolean unroll-factor=128' --cse --canonicalize main-secret-unroll-if-select.mlir > main-secret-bool.mlir

heir-opt --secret-distribute-generic --comb-to-cggi --cse --canonicalize main-secret-bool.mlir > main-secret-bool-cggi.mlir

heir-opt --cggi-to-tfhe-rust-bool --cse --canonicalize main-secret-bool-cggi.mlir > tfhe-bool.mlir

• The remove dead values call was me trying to see if it would remove the dead function.
• The entry function is called _Z5isqrti because polygeist outputs the mangled cpp name. We could fix this by adding extern "C" to the isqrt method. Noting this as the mangling might be different if you are using a different compiler.

Heir-opt output

The output has been copied into this gist since it is quite a lot of code.

After all of those passes, we get the linked mlir file which is supposed to be in the tfhe-bool-rust dialect.

The non-entry point method however has not been converted to tfhe-rust-bool dialect or removed. This might not be an issue, if we assume all function calls should be inlined -- I have not looked into the codebase to be able to deduce the correct behaviour.

Heir-translate

I then call heir-translate with:

heir-translate --emit-tfhe-rust-bool --use-levels tfhe-bool.mlir > circuit.rs

The output is also large and has been copied into this gist

The errors that occur in the first few lines are referring to the non-entrypoint function that has not been booleanized. Looking closer, polygeist has already inlined this function in the entrypoint, so this function should not be needed at all.

[j2kun]

As a potential workaround, the LLVM linkage attribute is what allows functions to be cleaned up as dead code

{llvm.linkage = #llvm.linkage<private>} instead of external, IIRC. But this is something we should fix.

We could have the emitter skip the functions it can't handle, or specify with a flag which function to codegen.

[j2kun]

This also reminds me that we should extend the pass that adds client helper encrypt/decrypt functions to support the rust codegen path.

[kevaundray]

As a potential workaround, the LLVM linkage attribute is what allows functions to be cleaned up as dead code

{llvm.linkage = #llvm.linkage<private>} instead of external, IIRC. But this is something we should fix.

We could have the emitter skip the functions it can't handle, or specify with a flag which function to codegen.

Thanks for the workaround :) TIL re the linkage attribute!