test_llvmstr.py

import numba                                                                                                                                                                                                                          
import math                                                                                                                                                                                                                           
import numpy as np                                                                                                                                                                                                                    
from numba import jit, float32, cfunc, int32, njit                                                                                                                                                                                    
                                                                                                                                                                                                                                      
                                                                                                                                                                                                                                      
import llvmlite.binding as llvm                                                                                                                                                                                                       
from ctypes.util import find_library                                                                                                                                                                                                  
from ctypes import (                                                                                                                                                                                                                  
        CFUNCTYPE, c_void_p, c_int, c_uint, c_ssize_t, c_ubyte, py_object, c_float,                                                                                                                                                   
        POINTER, byref,                                                                                                                                                                                                               
)                                                                                                                                                                                                                                     
                                                                                                                                                                                                                                      
def get_llvmir(func):                                                                                                                                                                                                                 
    llvm_obj = func.inspect_llvm()                                                                                                                                                                                                    
    if isinstance(llvm_obj, str):                                                                                                                                                                                                     
        llvm_ir = llvm_obj                                                                                                                                                                                                            
    else:                                                                                                                                                                                                                             
        for k,v in llvm_obj.items():                                                                                                                                                                                                  
            llvm_ir = v                                                                                                                                                                                                               
            break                                                                                                                                                                                                                     
    return llvm_ir                                                                                                                                                                                                                    
                                                                                                                                                                                                                                      
                                                                                                                                                                                                                                      
def create_execution_engine():                                                                                                                                                                                                        
    """                                                                                                                                                                                                                               
    Create an ExecutionEngine suitable for JIT code generation on                                                                                                                                                                     
    the host CPU.  The engine is reusable for an arbitrary number of                                                                                                                                                                  
    modules.                                                                                                                                                                                                                          
    """                                                                                                                                                                                                                               
    # Create a target machine representing the host                                                                                                                                                                                   
    target = llvm.Target.from_default_triple()                                                                                                                                                                                        
    target_machine = target.create_target_machine()                                                                                                                                                                                   
    # And an execution engine with an empty backing module                                                                                                                                                                            
    backing_mod = llvm.parse_assembly("")                                                                                                                                                                                             
    engine = llvm.create_mcjit_compiler(backing_mod, target_machine)                                                                                                                                                                  
    return engine                                                                                                                                                                                                                     
                                                                                                                                                                                                                                      
                                                                                                                                                                                                                                      
def compile_ir(engine, llvm_ir):                                                                                                                                                                                                      
    """                                                                                                                                                                                                                               
    Compile the LLVM IR string with the given engine.                                                                                                                                                                                 
    The compiled module object is returned.                                                                                                                                                                                           
    """                                                                                                                                                                                                                               
    # Create a LLVM module object from the IR                                                                                                                                                                                         
    mod = llvm.parse_assembly(llvm_ir)                                                                                                                                                                                                
    mod.verify()                                                                                                                                                                                                                      
    func_name = None                                                                                                                                                                                                                  
    for fid, func in enumerate(mod.functions):                                                                                                                                                                                        
                                                                                                                                                                                                  
        if func.name.startswith('cfunc'):                                                                                                                                                                                             
            func_name = func.name                                                                                                                                                                                                     
            break                                                                                                                                                                                                                     
    engine.add_module(mod)                                                                                                                                                                                                            
    engine.finalize_object()                                                                                                                                                                                                          
    engine.run_static_constructors()                                                                                                                                                                                                  
    return mod, func_name
    
@njit                                                                                                                                                                                                                                 
def udft(x, y):
    return len(x+'_'+y)
                                                                                                                                                                                             
                                                                                                                                                                                                                                      
a = 'ss'                                                                                                                                                                                                                              
b = 'a'                                                                                                                                                                                                                               
_c = udft(a, b)                                                                                                                                                                                                                       
print('init c:', _c)                                                                                                                                                                                                                  
                                                                                                                                                                                                                                      
llvm_ir = get_llvmir(udft)                                                                                                                                                                                                            
                                                                                                                                                                                                                                     
                                                                                                                                                                                                                
                                                                                                                                                                                                                                      
llvm.initialize()                                                                                                                                                                                                                     
llvm.initialize_native_target()                                                                                                                                                                                                       
llvm.initialize_native_asmprinter()  # yes, even this one                                                                                                                                                                             
                                                                                                                                                                                                                                      
engine = create_execution_engine()                                                                                                                                                                                                    
mod, func_name = compile_ir(engine, llvm_ir)                                                                                                                                                                                          
                                                                                                                                                                                                 
print('func_name:', func_name)                                                                                                                                                                                                        
                                                                                                                                                                                                                                      
func_ptr = engine.get_function_address(func_name)                                                                                                                                                                                     
print(func_ptr)                                                                                                                                                                                                                       
                                                                                                                                                                                                                                      
cfunc = CFUNCTYPE(c_int, py_object, py_object)(func_ptr)                                                                                                                                                                              
print('cfunc:', cfunc)                                                                                                                                                                                                                
                                                                                                                                                                                                                     
                                                                                                                                                                                                                
res = cfunc(a, b)