from __future__ import print_function import sys, os import ROOT os.environ['NUMBA_DEBUGINFO'] = '1' os.environ['NUMBA_DEVELOPER_MODE'] = '1' os.environ['NUMBA_DUMP_IR'] = '1' os.environ['NUMBA_DUMP_OPTIMIZED'] = '0' os.environ['NUMBA_DUMP_ASSEMBLY'] = '0' import root_utils from type_db import type_db from numba import types from numba import cgutils from numba.typing.typeof import typeof_impl from numba.typing.templates import signature from numba.targets.imputils import lower_builtin from llvmlite.llvmpy.core import Type #------------------------------------------------ # setup, workarounds, etc #------------------------------------------------ # I have problems calling fuctions returing void. It seems numba produces # ir like : $0.3 = call func($const0.2, k..... I.e. it ways assigs the return # value to something. For a void function the value is None, and we get # a failure in lowering.py, line 967 # Not sure why this happens, but if we type the function as returning an int, # the problem goes away. treat_void_return_as_int = True if treat_void_return_as_int: type_db.get(root_rep = 'void').llvm_type = Type.int(32) class RootMethodProxyType( types.Callable ): ''' numba type for all root methods The value will be a pointer to the 'this' object. (for static methods and global function, I have to see what to do (null ptr or a different type). ''' def __init__(self, arg = None ): if isinstance( arg, ROOT.MethodProxy ) : # MethodProxy.im_self is None if the method is not bound # or a regular function if not arg.im_self: self.tclass = None else : self.tclass = root_utils.get_tclass( arg.im_class ) name = arg.func_name elif isinstance( arg, ROOT.TMethod ): self.tclass = arg.GetClass() name = arg.GetName() print ("RootMethodProxyType ctor with arg = ", arg , " self.tclass",self.tclass ) self.methodname = name n = "RootMethodProxyType_"+name super(RootMethodProxyType, self).__init__(name=n) def get_call_type(self, context, args, kws): print ("Resolving call type for method", self.methodname, " of class", self.tclass, " with arguments", args ) # convert tuple of numba types args to string of arguments str_args = ",".join( [ type_db.get( numba_type = atype ).root_rep for atype in args ] ) # Find the correct method. ROOT is allowing for type conversions # so the method prototype does not need to correspond with # 'args' exactly. For the moment we can only hope that numba # will be able to cary out the required type conversions. if self.tclass : m = self.tclass.GetMethodWithPrototype( self.methodname, str_args ) else : m = ROOT.gROOT.GetGlobalFunctionWithPrototype ( self.methodname , str_args ) if not m : s = "Failed to find method of {} with name {} and arguments {} ".format ( self.tclass, self.methodname, str_args ) raise AttributeError( s ) else : print ("Method found", m, m.GetMangledName() , "retype", m.GetReturnTypeNormalizedName() ) retype = type_db.get( root_rep = m.GetReturnTypeNormalizedName() ).numba_type if treat_void_return_as_int : if retype == types.void : retype = types.int32 # We are going to return the actual function prototype, which we # get from root. def numba_type_from_methodarg ( marg ): typename = marg.GetTypeNormalizedName() return type_db.get( root_rep = typename ).numba_type arguments = map( numba_type_from_methodarg , m.GetListOfMethodArgs() ) r = signature( retype , *arguments ) if self.tclass : # It looks like setting recvr will put *this* in the first argument! r.recvr = self else : r.recvr = None print ("FFFFFFF", r, r.recvr ," retype = ", retype) print ( type(r), id(r) ) #------------------------------------------------------------------- # Define the lowering function here, while we still # hava access to the tmethod/tfunction (m) @lower_builtin( self, types.VarArg(types.Any) ) def impl_call(context, builder, sig, args): # does it matter if we use the same function name here?? print (" *************************** methodproxy lowering ***************************** ") print ( self ) print ( context ) print ( builder ) # todo: add some mechanism to avoid adding func many times F = root_utils.add_function( builder.basic_block.function.module , m ) print ("FFFFFF", F , sig, sig.return_type ) if sig.return_type != types.void : print ("xxxxxxxxxxxxxxxxxxxxxxxxxxxx") return builder.call( F , args ) else : rx = builder.call( F , args ) print ("xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx",rx ) # finally, return the signature... return r # we should not get here. raise TypeError("failed in resolve call type for method"+ str(self.methodname) + " of class" + str(self.tclass) + " with arguments", str(args) ) def get_call_signatures(self) : print (" get_call_signatures called") return (), True #================================================================== class RootObjectProxyType(types.Type): ''' The numba type for all root objectproxy objects. ''' def __init__(self, arg = None ): """ It is very important to make sure that the name given to the base ctor is unique for the tclass. There is some caching mechanism that relies on this. (uses 'key' property which defaults to name). """ tclass = None if isinstance( arg, ROOT.TClass ) : tclass = arg elif isinstance( arg, ROOT.ObjectProxy ): tclass = root_utils.get_tclass( arg ) if not tclass : raise TypeError("not a pyroot objectproxy or TClass", arg) name = 'RootObjectProxyType_' + tclass.GetName() self.tclass = tclass super(RootObjectProxyType, self).__init__(name=name) def __repr__( self ) : return "RootObjectProxyType for "+ str(self.tclass) def __str__( self ) : return self.__repr__() #----------------------------------------------------------------------------------------- # Register typers and objectmodels #----------------------------------------------------------------------------------------- from numba.extending import models, register_model @typeof_impl.register( ROOT.ObjectProxy ) def object_proxy_typer( val, c ) : return RootObjectProxyType( val ) @typeof_impl.register( ROOT.MethodProxy ) def method_proxy_typer( val, c ) : return RootMethodProxyType( val ) @register_model(RootObjectProxyType ) @register_model(RootMethodProxyType ) class RootMethodProxyTypeModel(models.OpaqueModel): # don't know what's opaque about it, but this seems def __init__(self, dmm, fe_type): # to be the model of choice -- also used for e.g. python objects models.OpaqueModel.__init__(self, dmm, fe_type) from numba.extending import unbox, NativeValue @unbox(RootObjectProxyType) def unbox_objectproxy(typ, obj, c): """ Convert object to a native structure ObjectProxy meth_proxy ---> +-------------------+ | pyobject_head | +---> CPP object +-------------------+ | | fObject | --+ | ... | +-------------------+ """ a = cgutils.pointer_add( c.builder, obj, root_utils.pyhead_size() , Type.pointer( Type.pointer( Type.int(8) ) ) ) a = c.builder.load( a ) return NativeValue( a ) @unbox(RootMethodProxyType) def unbox_methodproxy(typ, meth_proxy , c): """ MethodProxy ObjectProxy meth_proxy ---> +-------------------+ + ---> +-------------------+ | pyobject_head | | | pyobject_head | +---> CPP object +-------------------+ | +-------------------+ | | fSelf | --+ | fObject | --+ | ... | | ... | +-------------------+ +-------------------+ """ fSelf = cgutils.pointer_add( c.builder, meth_proxy, root_utils.pyhead_size() , Type.pointer( Type.pointer( Type.int(8) ) ) ) obj_proxy = c.builder.load( fSelf ) fObject = cgutils.pointer_add( c.builder, obj_proxy , root_utils.pyhead_size() , Type.pointer( Type.pointer( Type.int(8) ) )) r = c.builder.load( fObject ) return NativeValue( r ) from numba.typing.templates import (AttributeTemplate, infer_getattr ) @infer_getattr class RootProxyAttribute(AttributeTemplate): key = RootObjectProxyType def generic_resolve(self, record, attr): # We use root rtti to find the requested attribute #--------------------------- # it may be a a data-member #--------------------------- print("resolving attribute", attr, "for", record ) # todo: check whether this works on data members in base class dmember = record.tclass.GetDataMember( attr ) if dmember : print ("pointer? :", dmember.IsaPointer() ) # it seems getdatatype returns something for basic types if dmember.GetDataType() : t = type_db.get ( dmember.GetDataType() ) print (" typed " , t.numba_type ) return t.numba_type else : # we have some class # I use GetTypeName, which return the class name, even if the # datamember is a pointer. In the lowering, we have to deal with # this. tclass = ROOT.TClass.GetClass( dmember.GetTypeName(), False, False ) if not tclass : raise AttributeError("failed to get tclass for type" + str( dmember.GetTypeName() ) ) RT = RootObjectProxyType( tclass ) print (" typed " , RT ) return RT #-------------------- # it may be a method #-------------------- method = record.tclass.GetMethodAllAny( attr ) if method : RT = RootMethodProxyType( method ) print ( "typed" , RT , method ) return RT # failed to resolve the attribute return None from numba.targets.imputils import lower_getattr_generic from llvmlite import ir @lower_getattr_generic( RootObjectProxyType ) def interval_getattr(context, builder, typ, value, attr): """ Generic getattr() implementation """ print ('generic gettattr implementation for', typ.tclass) print ('attr=', attr) dmember = typ.tclass.GetDataMember( attr ) if dmember : if dmember.GetDataType() : t = type_db.get ( dmember.GetDataType() ) # add the offset and interpret as pointer to the correct type offset = dmember.GetOffset() pval = cgutils.pointer_add( builder, value , offset, return_type=Type.pointer( t.llvm_type ) ) pval = builder.load( pval ) return pval else : # construct a new RootObjectProxyType and point it to the attribute tclass = ROOT.TClass.GetClass( dmember.GetTypeName(), False, False ) RT = RootObjectProxyType( tclass ) #numba_obj = cgutils.create_struct_proxy(RT)( context, builder) offset = dmember.GetOffset() print ("offset = ", offset ) # if the member is value, simply set addres in the new proxy, # if it is a pointer we need to deref it if dmember.IsaPointer() : addr = cgutils.pointer_add( builder, value , offset , Type.pointer( Type.pointer( Type.int(8) ))) cpp_addr = builder.load( addr ) else : addr = cgutils.pointer_add( builder, value , offset ) cpp_addr = addr return cpp_addr._getvalue() method = typ.tclass.GetMethodAllAny( attr ) if method: # The value of a RootMethodProxyType object is a pointer to its # *this* object. return value print ("Failed to lower getattr ") # should not get here res = context.get_constant( types.int32, 43 ) return res #==================================================================================================================== #==================================================================================================================== #-------------------------------------------------------------- # Load c++ code for our classes into cling #-------------------------------------------------------------- ROOT.gInterpreter.Declare( open("cpp_for_testing.cc").read() ) aa = ROOT.AA() print ( aa.funkd(100) ) from numba import jit @jit(nopython=True) def bee(obj): print ('result of obj.a_int =', obj.a_int) print (' a_unsigned_int =', obj.a_unsigned_int ) print (' a_float =', obj.a_float) print (' a_double =', obj.a_double) print (' a_bool =', obj.a_bool) return 420000 @jit(nopython=True) def testbb( bb ) : x = bb.aa print ( x.a_int ) print ( bb.aa_ptr.a_int ) print ( bb.int_ptr ) return 124 @jit(nopython=True) def testmethod( m ) : m(1) return 0 @jit(nopython=True) def testmethod2( aa ): aa.funkd( 1.1 ) testmethod( ROOT.a_void_function ) testmethod.inspect_types() if False : #m = ROOT.TH1D.Fill # any methodproxy #h = ROOT.TH1D("h","h",10,0,100) #m = h.Fill #testmethod(m) aa = ROOT.AA() aa.a_int = 12345 testmethod2( aa ) aa.a_int = -42 aa.a_unsigned_int = 42 aa.a_float = 42.0 aa.a_double = 4242.0 aa.a_bool = True bb = ROOT.BB() bb.aa.a_int = -123 bb.aa_ptr = aa # pyroot does this right bb.look() print ("int ptr", bb.int_ptr ) print (testbb(bb))