Help needed wih deferred types and StructRef

Numba Support: How do I do ...?
1

While I was trying to use the deferred_type with StructRefs I get a SIGABRT in the following example:


import numba as nb
from numba import njit, types, typed, float64, int64
from numba.core import types
from numba.experimental import structref
from numba.core.extending import overload_method

@structref.register
class VectorStructType(types.StructRef):
    def preprocess_fields(self, fields):
        return tuple((name, types.unliteral(typ)) for name, typ in fields)

class VectorStruct(structref.StructRefProxy):
    def __new__(cls, x, y, z):
        return structref.StructRefProxy.__new__(cls, x, y, z)

    @property
    def x(self):
        return VectorStruct_get_x(self)

    @property
    def y(self):
        return VectorStruct_get_y(self)

    @property
    def z(self):
        return VectorStruct_get_z(self)

@njit
def VectorStruct_get_x(self):
    return self.x

@njit
def VectorStruct_get_y(self):
    return self.y

@njit
def VectorStruct_get_z(self):
    return self.z

structref.define_proxy(VectorStruct, VectorStructType, ["x", "y", "z"])

vector_struct_type = VectorStructType(fields=(('x', float64), ('y', float64),
                                              ('z', float64)))

PolygonTypeDeferred = types.deferred_type()
@structref.register
class PolygonStructType(types.StructRef):
    def preprocess_fields(self, fields):
        return tuple([
            ('vertices', types.Optional(types.ListType(vector_struct_type))),
            ('parent', types.Optional(PolygonTypeDeferred)),
            ])

PolygonType_inst = PolygonStructType(fields=(('vertices', types.Optional(types.ListType(vector_struct_type))), ('parent', types.Optional(PolygonTypeDeferred))))
PolygonTypeDeferred.define(PolygonType_inst)

class PolygonStruct(structref.StructRefProxy):
    def __new__(cls, vertices, parent):
        return structref.StructRefProxy.__new__(cls, vertices, parent)

    @property
    def vertices(self):
        return PolygonStruct_get_vertices(self)

    @property
    def parent(self):
        return PolygonStruct_get_parent(self)

@njit
def PolygonStruct_get_vertices(self):
    return self.vertices

@njit
def PolygonStruct_get_parent(self):
    return self.parent

structref.define_proxy(PolygonStruct, PolygonStructType, ["vertices", "parent"])

@overload_method(PolygonStructType, "flip")
def PolygonStruct_flip(self):
    def impl(self):
        if self.vertices is not None:
            ref = self.vertices[0]
            ref.x = -ref.x
    return impl

@njit
def test():
    lst = typed.List.empty_list(vector_struct_type)
    lst.append(VectorStruct(1.0, 2.0, 3.0))
    lst[0].x = 5.0

    print('before')
    p = PolygonStruct(None, None)
    p.flip()
    p.vertices = lst
    p.flip()
    print('after')
    return p

print('start')
y = test()
p2 = PolygonStruct(None, None)
print(y.vertices[0].x)
print('end')

Ticket open: https://github.com/numba/numba/issues/6077

2

The workaround I found so far, and it appears to be working, is to use jitclass with init only instead of StructRef and overload_method decorator. Like this:

from numba import typed, typeof, njit, int64, types
from numba.experimental import jitclass
from numba.extending import overload_method, overload

NodeType = types.deferred_type()
@jitclass([('value', int64), ('parent', types.Optional(NodeType))])
class Node(object):
    def __init__(self, value, parent):
        self.value = value
        self.parent = parent

NodeType.define(Node.class_type.instance_type)   

@overload_method(types.misc.ClassInstanceType, 'test')
def test_node(inst,):
    if inst is Node.class_type.instance_type:
        def impl(inst,):
            if inst.parent is not None:
                print('has parent')
            else:
                print('has no parent')
        return impl

@njit
def test():
    n = Node(5, None)
    n.test()
    p = Node(0, None)
    n.parent = p
    n.test()

test()
1 Like