How to call into LLVM-IR code that has multiple functions and global variables

Numba Community Support
1

I just wanted to try to call multipe functions from LLMIR, which interact through a global variable (in this case the function pointer of this scipy.special function which has to be initiated at runtime). This is just an example, but should also result in an easy way of enable caching and potentially adding more functions, which are not easy to be called by Numba (eg. complex numbers involved).

How is it possible to call from multiple Numba functions into the same LLVM-IR module?

import numba as nb
from numba import njit,types
from llvmlite import ir, binding as ll
from llvmlite import binding

cython_special = r"""
@cy_voigt_profile = dso_local local_unnamed_addr global double (double, double, double, i32)* null, align 8

; Function Attrs: mustprogress nofree norecurse nosync nounwind uwtable willreturn writeonly
define dso_local void @init_cy_voigt_profile(i64 noundef %0) local_unnamed_addr #0 {
  %2 = inttoptr i64 %0 to double (double, double, double, i32)*
  store double (double, double, double, i32)* %2, double (double, double, double, i32)** @cy_voigt_profile, align 8, !tbaa !4
  ret void
}

; Function Attrs: mustprogress nofree norecurse nosync nounwind readonly uwtable willreturn
define dso_local i64 @adress_cy_voigt_profile() local_unnamed_addr #2 {
  %1 = load double (double, double, double, i32)*, double (double, double, double, i32)** @cy_voigt_profile, align 8, !tbaa !4
  %2 = ptrtoint double (double, double, double, i32)* %1 to i64
  ret i64 %2
}

attributes #0 = { mustprogress nofree norecurse nosync nounwind uwtable willreturn writeonly "frame-pointer"="none" "min-legal-vector-width"="0" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "tune-cpu"="generic" }
attributes #1 = { nounwind uwtable "frame-pointer"="none" "min-legal-vector-width"="0" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "tune-cpu"="generic" }
attributes #2 = { mustprogress nofree norecurse nosync nounwind readonly uwtable willreturn "frame-pointer"="none" "min-legal-vector-width"="0" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "tune-cpu"="generic" }
attributes #3 = { nounwind }

!llvm.module.flags = !{!0, !1, !2}
!llvm.ident = !{!3}

!0 = !{i32 1, !"wchar_size", i32 2}
!1 = !{i32 7, !"PIC Level", i32 2}
!2 = !{i32 7, !"uwtable", i32 1}
!3 = !{!"clang version 14.0.6"}
!4 = !{!5, !5, i64 0}
!5 = !{!"any pointer", !6, i64 0}
!6 = !{!"omnipotent char", !7, i64 0}
!7 = !{!"Simple C/C++ TBAA"}
"""

@nb.extending.intrinsic
def init_cy_voigt_profile(typingctx, args):
    def compile_library(context, asm):
        library = context.codegen().create_library('scipy_special')
        ll_module = ll.parse_assembly(asm)
        ll_module.verify()
        library.add_llvm_module(ll_module)
        return library

    def codegen(context, builder, sig, args):
        library = compile_library(context, cython_special)
        context.active_code_library.add_linking_library(library)
        argtypes = [context.get_argument_type(aty) for aty in sig.args]
        restype = context.get_argument_type(sig.return_type)
        fnty = ir.FunctionType(restype, argtypes)
        fn = nb.core.cgutils.insert_pure_function(builder.module, fnty, name="init_cy_voigt_profile")
        retval = context.call_external_function(builder, fn, sig.args, args)
        return retval

    sig = types.int64(types.int64)
    return sig, codegen

@nb.extending.intrinsic
def adress_cy_voigt_profile(typingctx):
    def compile_library(context, asm):
        library = context.codegen().create_library('scipy_special')
        ll_module = ll.parse_assembly(asm)
        ll_module.verify()
        library.add_llvm_module(ll_module)
        return library

    def codegen(context, builder, sig, args):
        library = compile_library(context, cython_special)
        context.active_code_library.add_linking_library(library)
        argtypes = [context.get_argument_type(aty) for aty in sig.args]
        restype = context.get_argument_type(sig.return_type)
        fnty = ir.FunctionType(restype, argtypes)
        fn = nb.core.cgutils.insert_pure_function(builder.module, fnty, name="adress_cy_voigt_profile")
        retval = context.call_external_function(builder, fn, sig.args, args)
        return retval

    sig = types.int64()
    return sig, codegen


@njit()
def init(adress):
    init_cy_voigt_profile(adress)

@njit()
def adress_voigt_profile():
    return adress_cy_voigt_profile()

out = init(125)
#obviously wrong
adress_voigt_profile()
2

Would this help? I tested with ctypes only, but hopefully should work in numba context as well.

from ctypes import CFUNCTYPE, c_int64
from llvmlite import binding as ll


cython_special = ...  # as in the OP


ll.initialize()
ll.initialize_native_target()
ll.initialize_native_asmprinter()
llvm_cython_special = ll.parse_assembly(cython_special)
llvm_cython_special.verify()
target_machine = ll.Target.from_default_triple().create_target_machine()
engine_cython_special = ll.create_mcjit_compiler(llvm_cython_special, target_machine)
engine_cython_special.finalize_object()

init_p = engine_cython_special.get_function_address("init_cy_voigt_profile")
address_p = engine_cython_special.get_function_address("adress_cy_voigt_profile")


init_fn = CFUNCTYPE(None, c_int64)(init_p)
address_fn = CFUNCTYPE(c_int64)(address_p)


if __name__ == '__main__':
    a = 123
    init_fn(a)
    assert address_fn() == a
1 Like