[OCaml] Add LLVMInternalizePredicateBindings

[kren1]

Summary

This patch exposes the predicate API of internalize pass to OCaml.

[whitequark]

bindings/ocaml/transforms/ipo/llvm_ipo.mli

76

I don't think we make any claims of thread safety for any other API. Why a warning on this one?

[kren1]

bindings/ocaml/transforms/ipo/llvm_ipo.mli

76

I felt it's necessary to warn the user that the predicate is not passed by value, but its reference is stored in a global variable. This kind of seemed like the most concise way of saying that. Would you prefer it rephrased to more closely match the above or just removed?

[whitequark]

bindings/ocaml/transforms/ipo/llvm_ipo.mli

76

Wait. That's actually worse, because it's not *non-thread-safe*, it's not *reentrant*. Which probably doesn't matter for this specific function, but would be good to solve in general.

I haven't written any really complex OCaml/C interop code in a few years, so my skills here are a bit rusty... could you please look into a better way to do this? If not, no worries, I'll refresh my memory and dig into it myself, but it might take some weeks to get around to this.

[kren1]

Make add_internalize_predicate reentrant

[kren1]

bindings/ocaml/transforms/ipo/llvm_ipo.mli

76

I think I addressed this now. Cleaning up is a bit ugly, because I think the lifetime of the predicate closure should basically match the lifetime of the PassManager. I found no examples where this would be handled elsewhere. So I implemented a global association of PassManager and Predicate closure. When you dispose of the pass manager you walk the association and free all predicated associated with the pass manager being disposed.

Another option would be to extend the PassManager API and make it know about the closures associated with it, but I think this binding has little to do with the PassManager and should stay in its own directory.

[whitequark]

bindings/ocaml/transforms/ipo/llvm_ipo.mli

76

Thank you. Correct me if I'm wrong, but it looks like you could use OCaml APIs to maintain a linked list of predicates, and doing so would simplify the memory management here quite a bit, as well as make it more obviously sound.

[kren1]

bindings/ocaml/transforms/ipo/llvm_ipo.mli

76

You mean use caml_alloc and Store_field to construct an OCaml linked list in C? I could do that, but not sure it would be cleaner and more obviously sound. I find it easier to reason about C than what OCaml GC will be doing.

Another approach could be do to some mix-and-match of OCaml and C, for example deal with allocating and traversing the list in OCaml and then have C only allocate the contents of each list. That might work, but it would be spread over atleast 5 files and would require the Llvm_ipo to have a public ref exposed in its api (so that the PassManager can access it).

There is a similar problem in the first approach but there I could have an extern value* (similar to how I do it now).

Another way would be to find an implementation of linked-list in OCaml/LLVM C-api and then use that. That would be my preferred approach, but I couldn't find anything useable there (OCaml seems very GC orientated and LLVM has nothing). Do you have any pointers?

Let me know what would be the best way to proceed.

[kren1]

Move most of the list handling stuff to OCaml

[kren1]

Fix the List.filter predicate

bindings/ocaml/llvm/llvm_ocaml.c

CAMLprim value llvm_passmanager_run_function(LLVMValueRef F,
   return Val_bool(LLVMRunFunctionPassManager(FPM, F));
 }
 
 /* [ `Function ] PassManager.t -> bool */
 CAMLprim value llvm_passmanager_finalize(LLVMPassManagerRef FPM) {
   return Val_bool(LLVMFinalizeFunctionPassManager(FPM));
 }
 
+/* (LLVM.PassManager.t, (string -> bool)) list
+   Holds all the closures that */
+extern value *active_closures_list;
+
 /* PassManager.any PassManager.t -> unit */
 CAMLprim value llvm_passmanager_dispose(LLVMPassManagerRef PM) {
+
+  // We drop any closures associated with this PassManager. The closures will
+  // then be GCed
+  static value *drop_closure_f = NULL;
+  if (drop_closure_f == NULL) {
+    drop_closure_f = caml_named_value("ll_drop_closure");
+  }
+  if (active_closures_list != NULL) {
+    *active_closures_list =
+        caml_callback2(*drop_closure_f, *active_closures_list, (value)PM);
+  }
+
   LLVMDisposePassManager(PM);
   return Val_unit;
 }

bindings/ocaml/transforms/ipo/ipo_ocaml.c

 |* This file glues LLVM's OCaml interface to its C interface. These functions *|
 |* are by and large transparent wrappers to the corresponding C functions.    *|
 |*                                                                            *|
 |* Note that these functions intentionally take liberties with the CAMLparamX *|
 |* macros, since most of the parameters are not GC heap objects.              *|
 |*                                                                            *|
 \*===----------------------------------------------------------------------===*/
 
-#include "llvm-c/Transforms/IPO.h"
-#include "caml/mlvalues.h"
+#include "caml/alloc.h"
+#include "caml/callback.h"
+#include "caml/fail.h"
+#include "caml/memory.h"
 #include "caml/misc.h"
+#include "caml/mlvalues.h"
+#include "llvm-c/Core.h"
+#include "llvm-c/Transforms/IPO.h"
+#include <assert.h>
 
 /* [`Module] Llvm.PassManager.t -> unit */
 CAMLprim value llvm_add_argument_promotion(LLVMPassManagerRef PM) {
   LLVMAddArgumentPromotionPass(PM);
   return Val_unit;
 }
 
 /* [`Module] Llvm.PassManager.t -> unit */

 }
 
 /* [`Module] Llvm.PassManager.t -> all_but_main:bool -> unit */
 CAMLprim value llvm_add_internalize(LLVMPassManagerRef PM, value AllButMain) {
   LLVMAddInternalizePass(PM, Bool_val(AllButMain));
   return Val_unit;
 }
 
+/* (LLVM.PassManager.t, (string -> bool)) list
+   Memory managed by OCaml GC*/
+value *active_closures_list = NULL;
+
+LLVMBool MustPreserveCallBack(LLVMValueRef Val, void *Ctx) {
+  CAMLparam0();
+  CAMLlocal2(LLVMValName, Predicate);
+  const char *llvmValName;
+  LLVMBool ret;
+  LLVMPassManagerRef PM = (LLVMPassManagerRef)Ctx;
+  static value *find_closure_f = NULL;
+  if (find_closure_f == NULL) {
+    find_closure_f = caml_named_value("ll_find_closure");
+  }
+  Predicate = caml_callback2(*find_closure_f, *active_closures_list, (value)PM);
+  llvmValName = LLVMGetValueName(Val);
+  LLVMValName = caml_copy_string(llvmValName);
+  ret = Bool_val(caml_callback(Predicate, LLVMValName));
+
+  CAMLreturnT(LLVMBool, ret);
+}
+
+/* [`Module] Llvm.PassManager.t -> (string -> bool) -> unit */
+CAMLprim value llvm_add_internalize_predicate(LLVMPassManagerRef PM,
+                                              value Predicate) {
+  CAMLparam1(Predicate);
+  if (active_closures_list == NULL) {
+    // initialize the active_closure_list
+    active_closures_list = malloc(sizeof(value));
+    if (active_closures_list == NULL)
+      caml_raise_out_of_memory();
+    *active_closures_list = Val_emptylist;
+    caml_register_global_root(active_closures_list);
+  }
+
+  static value *add_or_replace_closure_f = NULL;
+  if (add_or_replace_closure_f == NULL) {
+    add_or_replace_closure_f = caml_named_value("ll_add_or_replace_closure");
+  }
+  *active_closures_list = caml_callback3(
+      *add_or_replace_closure_f, *active_closures_list, (value)PM, Predicate);
+
+  /* We can't directly pass the Predicate value to the callback as GC might move
+     it
+     before callback calls it. So we pass it the PM pointer, which is not GC
+     dependant. */
+  LLVMAddInternalizePassWithMustPreservePredicate(PM, PM,
+                                                  &MustPreserveCallBack);
+  CAMLreturn(Val_unit);
+}
+
 /* [`Module] Llvm.PassManager.t -> unit */
 CAMLprim value llvm_add_strip_dead_prototypes(LLVMPassManagerRef PM) {
   LLVMAddStripDeadPrototypesPass(PM);
   return Val_unit;
 }
 
 /* [`Module] Llvm.PassManager.t -> unit */
 CAMLprim value llvm_add_strip_symbols(LLVMPassManagerRef PM) {
   LLVMAddStripSymbolsPass(PM);
   return Val_unit;
 }

bindings/ocaml/transforms/ipo/llvm_ipo.ml

external add_prune_eh
   : [ `Module ] Llvm.PassManager.t -> unit
   = "llvm_add_prune_eh"
 external add_ipsccp
   : [ `Module ] Llvm.PassManager.t -> unit
   = "llvm_add_ipsccp"
 external add_internalize
   : [ `Module ] Llvm.PassManager.t -> all_but_main:bool -> unit
   = "llvm_add_internalize"
+external add_internalize_predicate
+  : [ `Module ] Llvm.PassManager.t -> (string -> bool) -> unit
+  = "llvm_add_internalize_predicate"
 external add_strip_dead_prototypes
   : [ `Module ] Llvm.PassManager.t -> unit
   = "llvm_add_strip_dead_prototypes"
 external add_strip_symbols
   : [ `Module ] Llvm.PassManager.t -> unit
   = "llvm_add_strip_symbols"
+
+
+(* Helper functions for add_internalize_predicate. Used by the C code *)
+let add_or_replace_closure :
+       ([`Module] Llvm.PassManager.t * (string -> bool)) list
+    -> [`Module] Llvm.PassManager.t
+    -> (string -> bool)
+    -> ([`Module] Llvm.PassManager.t * (string -> bool)) list =
+ fun ls pm c -> (pm, c) :: List.filter (fun (pm', _) -> pm' == pm) ls
+
+let rec find_closure :
+       ([`Module] Llvm.PassManager.t * (string -> bool)) list
+    -> [`Module] Llvm.PassManager.t
+    -> string
+    -> bool =
+ fun ls pm ->
+  match ls with
+  | [] -> assert false
+  | (pm', c) :: _ when pm' == pm -> c
+  | head :: tail -> find_closure tail pm
+
+let drop_closure :
+       ([`Module] Llvm.PassManager.t * (string -> bool)) list
+    -> [`Module] Llvm.PassManager.t
+    -> ([`Module] Llvm.PassManager.t * (string -> bool)) list =
+ fun ls pm -> List.filter (fun (pm', _) -> pm' == pm) ls
+
+let _ = Callback.register "ll_add_or_replace_closure" add_or_replace_closure
+let _ = Callback.register "ll_find_closure" find_closure
+let _ = Callback.register "ll_drop_closure" drop_closure

bindings/ocaml/transforms/ipo/llvm_ipo.mli

external add_ipsccp
   : [ `Module ] Llvm.PassManager.t -> unit
   = "llvm_add_ipsccp"
 
 (** See the [llvm::createInternalizePass] function. *)
 external add_internalize
   : [ `Module ] Llvm.PassManager.t -> all_but_main:bool -> unit
   = "llvm_add_internalize"
 
+(** See the [llvm::createInternalizePass] function.
+    If predicate returns [true], that symbol is preserved. *)
+val add_internalize_predicate
+  : [ `Module ] Llvm.PassManager.t -> (string -> bool) -> unit
+
 (** See the [llvm::createStripDeadPrototypesPass] function. *)
 external add_strip_dead_prototypes
   : [ `Module ] Llvm.PassManager.t -> unit
   = "llvm_add_strip_dead_prototypes"
 
 (** See the [llvm::createStripSymbolsPass] function. *)
 external add_strip_symbols
   : [ `Module ] Llvm.PassManager.t -> unit
   = "llvm_add_strip_symbols"