[P1] Fix `leverage-predicated-value pass` for c++ lowering to neura

[tancheng]

Make --leverage-predicated-value work for bra. For now, a basic block in the target function is referred to by a branch operation from the main function, leading to type violation. This needs a proper design for this.

[ShangkunLi]

Why would this bring type violation? We can transform branch operations into dataflow-like operations (e.g., phi, ctrl_mov, etc.).

[tancheng]

Sorry, wrong description, PTAL at the test in #85, when --leverage-predicated-value pass is enabled, segmentation fault is encountered. Should be easy to fix by you~! I appreciate it in advance~!

module attributes {dlti.dl_spec = #dlti.dl_spec<!llvm.ptr = dense<64> : vector<4xi64>, i128 = dense<128> : vector<2xi64>, i1 = dense<8> : vector<2xi64>, i64 = dense<64> : vector<2xi64>, i32 = dense<32> : vector<2xi64>, f16 = dense<16> : vector<2xi64>, i16 = dense<[16, 32]> : vector<2xi64>, f64 = dense<64> : vector<2xi64>, i8 = dense<[8, 32]> : vector<2xi64>, f128 = dense<128> : vector<2xi64>, "dlti.stack_alignment" = 128 : i64, "dlti.endianness" = "little">, llvm.ident = "Ubuntu clang version 14.0.0-1ubuntu1.1"} {
  llvm.mlir.global external local_unnamed_addr @input(dense<1.000000e+00> : tensor<32xf32>) {addr_space = 0 : i32, alignment = 4 : i64, dso_local} : !llvm.array<32 x f32>
  llvm.mlir.global external local_unnamed_addr @output(dense<0.000000e+00> : tensor<32xf32>) {addr_space = 0 : i32, alignment = 4 : i64, dso_local} : !llvm.array<32 x f32>
  llvm.mlir.global external local_unnamed_addr @coefficients(dense<[2.500000e-01, 1.500000e+00, 3.750000e+00, -2.250000e+00, 5.000000e-01, 7.500000e-01, -3.000000e+00, 1.250000e+00, 2.500000e-01, 1.500000e+00, 3.750000e+00, -2.250000e+00, 5.000000e-01, 7.500000e-01, -3.000000e+00, 1.250000e+00, 2.500000e-01, 1.500000e+00, 3.750000e+00, -2.250000e+00, 5.000000e-01, 7.500000e-01, -3.000000e+00, 1.250000e+00, 2.500000e-01, 1.500000e+00, 3.750000e+00, -2.250000e+00, 5.000000e-01, 7.500000e-01, -3.000000e+00, 1.250000e+00]> : tensor<32xf32>) {addr_space = 0 : i32, alignment = 4 : i64, dso_local} : !llvm.array<32 x f32>
  llvm.mlir.global private unnamed_addr constant @".str"("output: %f\0A\00") {addr_space = 0 : i32, alignment = 1 : i64, dso_local}
  llvm.func local_unnamed_addr @main() -> (i32 {llvm.noundef}) attributes {frame_pointer = #llvm.framePointerKind<"non-leaf">, no_unwind, passthrough = ["mustprogress", "nofree", "norecurse", ["uwtable", "2"], ["min-legal-vector-width", "0"], ["no-trapping-math", "true"], ["stack-protector-buffer-size", "8"], ["target-cpu", "generic"]], target_cpu = "generic", target_features = #llvm.target_features<["+neon", "+outline-atomics", "+v8a"]>} {
    %0 = llvm.mlir.addressof @output : !llvm.ptr
    %1 = llvm.mlir.constant(0 : i64) : i64
    %2 = llvm.mlir.addressof @input : !llvm.ptr
    %3 = llvm.mlir.addressof @coefficients : !llvm.ptr
    %4 = llvm.mlir.constant(1 : i64) : i64
    %5 = llvm.mlir.constant(32 : i64) : i64
    %6 = llvm.mlir.addressof @".str" : !llvm.ptr
    %7 = llvm.mlir.constant(0 : i32) : i32
    %8 = llvm.load %0 {alignment = 4 : i64, tbaa = [#tbaa_tag]} : !llvm.ptr -> f32
    llvm.br ^bb1(%8, %1 : f32, i64)
  ^bb1(%9: f32, %10: i64):  // 2 preds: ^bb0, ^bb1
    %11 = llvm.getelementptr inbounds %2[0, %10] : (!llvm.ptr, i64) -> !llvm.ptr, !llvm.array<32 x f32>
    %12 = llvm.load %11 {alignment = 4 : i64, tbaa = [#tbaa_tag]} : !llvm.ptr -> f32
    %13 = llvm.getelementptr inbounds %3[0, %10] : (!llvm.ptr, i64) -> !llvm.ptr, !llvm.array<32 x f32>
    %14 = llvm.load %13 {alignment = 4 : i64, tbaa = [#tbaa_tag]} : !llvm.ptr -> f32
    %15 = llvm.fmul %12, %14 : f32
    %16 = llvm.fadd %9, %15 : f32
    %17 = llvm.add %10, %4 overflow<nsw, nuw> : i64
    %18 = llvm.icmp "eq" %17, %5 : i64
    llvm.cond_br %18, ^bb2, ^bb1(%16, %17 : f32, i64) {loop_annotation = #loop_annotation}
  ^bb2:  // pred: ^bb1
    llvm.store %16, %0 {alignment = 4 : i64, tbaa = [#tbaa_tag]} : f32, !llvm.ptr
    %19 = llvm.fpext %16 : f32 to f64
    %20 = llvm.call @printf(%6, %19) vararg(!llvm.func<i32 (ptr, ...)>) {no_unwind} : (!llvm.ptr, f64) -> i32
    llvm.return %7 : i32
  }
  llvm.func local_unnamed_addr @_Z6kernelPfS_S_(%arg0: !llvm.ptr {llvm.nocapture, llvm.noundef, llvm.readonly}, %arg1: !llvm.ptr {llvm.nocapture, llvm.noundef}, %arg2: !llvm.ptr {llvm.nocapture, llvm.noundef, llvm.readonly}) attributes {accelerator = "neura", frame_pointer = #llvm.framePointerKind<"non-leaf">, no_unwind, passthrough = ["mustprogress", "nofree", "norecurse", "nosync", ["uwtable", "2"], ["min-legal-vector-width", "0"], ["no-trapping-math", "true"], ["stack-protector-buffer-size", "8"], ["target-cpu", "generic"]], target_cpu = "generic", target_features = #llvm.target_features<["+neon", "+outline-atomics", "+v8a"]>} {
    %0 = "neura.constant"() <{predicate = true, value = 0 : i64}> : () -> i64
    %1 = "neura.constant"() <{predicate = true, value = 1 : i64}> : () -> i64
    %2 = "neura.constant"() <{predicate = true, value = 32 : i64}> : () -> i64
    neura.br %0 : i64 to ^bb1
  ^bb1(%3: i64):  // 2 preds: ^bb0, ^bb1
    %4 = "neura.gep"(%arg0, %3) : (!llvm.ptr, i64) -> !llvm.ptr
    %5 = "neura.load"(%4) : (!llvm.ptr) -> f32
    %6 = "neura.gep"(%arg2, %3) : (!llvm.ptr, i64) -> !llvm.ptr
    %7 = "neura.load"(%6) : (!llvm.ptr) -> f32
    %8 = "neura.load"(%arg1) : (!llvm.ptr) -> f32
    %9 = "neura.fmul_fadd"(%5, %7, %8) : (f32, f32, f32) -> f32
    "neura.store"(%9, %arg1) : (f32, !llvm.ptr) -> ()
    %10 = "neura.add"(%3, %1) : (i64, i64) -> i64
    %11 = "neura.icmp"(%10, %2) <{cmpType = "eq"}> : (i64, i64) -> i1
    neura.cond_br %11 : i1 then to ^bb2 else %10 : i64 to ^bb1
  ^bb2:  // pred: ^bb1
    "neura.return"() : () -> ()
  }
  llvm.func local_unnamed_addr @printf(!llvm.ptr {llvm.nocapture, llvm.noundef, llvm.readonly}, ...) -> (i32 {llvm.noundef}) attributes {frame_pointer = #llvm.framePointerKind<"non-leaf">, no_unwind, passthrough = ["nofree", ["no-trapping-math", "true"], ["stack-protector-buffer-size", "8"], ["target-cpu", "generic"]], target_cpu = "generic", target_features = #llvm.target_features<["+neon", "+outline-atomics", "+v8a"]>}
}

[ShangkunLi]

This problem can be fixed together with the TODO in #73.

I also encountered such segmentation faults before when using clang++ as the front end. It is probably because:

• We do not have an llvm.ptr operation in our neura dialect
• We do not transform the func arguments into the predicated type.

For now, I just use polygeist as the frontend (not clang++) because the affine ir generated avoids handling such llvm.ptr operations. And I can use load_indexed/store_indexed to avoid handling gep and ptr operations.

[tancheng]

We can also apply existing grant_predicate or grant_always onto the llvm.ptr, which sounds more robust and makes the target func be co-simulated with any llvm caller.

[ShangkunLi]

Fixed in #86.

[tancheng]

Excellent~! Thanks!