def test_xbuffer_operators(self): a = np.array([1., 1.5], dtype=np.float32) xb = XBuffer(a) xb1 = xb + [1, -1] assert isinstance(xb1, XBuffer) np.testing.assert_equal(xb.to_numpy(), a) np.testing.assert_equal(xb1.to_numpy(), np.array([2., .5], dtype=np.float32))
def test_xbuffer_arg_name(self): def py_func(xb): np.testing.assert_equal(xb.to_numpy(), np.array([1., 2.], dtype=np.float32)) xb.copy_from(xb * 2) xb = XBuffer(np.array([1., 2.], dtype=np.float32)) of = OpaqueFunc(py_func, [TypeCode.XBuffer]) of(xb) np.testing.assert_equal(xb.to_numpy(), np.array([2., 4.], dtype=np.float32))
def test_xbuffer_equals(self): a = np.array([1., 1.5], dtype=np.float32) xb = XBuffer(a) b = xb.to_numpy() xb *= np.array([2, -1]) assert isinstance(xb, XBuffer) np.testing.assert_equal(xb.to_numpy(), np.array([2., -1.5], dtype=np.float32)) np.testing.assert_equal(a, np.array([1., 1.5], dtype=np.float32)) np.testing.assert_equal(b, np.array([1., 1.5], dtype=np.float32))
def test_xbuffer_getattr(self): a = np.array([1., 1.5], dtype=np.float32) xb = XBuffer(a) assert xb.dtype == 'float32' assert xb.shape == (2, ) a = np.zeros((1, 3, 224, 224), dtype=np.int8) xb = XBuffer(a) assert xb.dtype == 'int8' assert xb.shape == (1, 3, 224, 224)
def test_build_rt_opaque_func_cpu_tf(self): tf.compat.v1.reset_default_graph() W = np.reshape( np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]], dtype=np.float32), (2, 1, 2, 2)) # B = np.array([0, 0], dtype=np.float32) iX = px.ops.input('in', [1, 1, 4, 4]) wX = px.ops.constant('w', W) # bX = px.ops.constant('b', B) cX = px.ops.conv2d('conv', iX, wX, kernel_size=[2, 2], strides=[1, 1], padding_hw=[0, 0, 0, 0], dilation=[1, 1], groups=1, channels=2, data_layout='NCHW', kernel_layout='OIHW') xlayers = [iX, wX, cX] xg = TestBase.xg_factory.build_from_xlayer(xlayers) of = OpaqueFuncRegistry.Get('pyxir.build_rt') rt_of = OpaqueFunc() # call opaque func of(xg, 'cpu', 'cpu-tf', ['in'], ['conv'], rt_of) # By now, the `rt_of` is initialized with the opaque runtime function ins = [XBuffer(np.ones((1, 1, 4, 4), dtype=np.float32))] outs = [XBuffer(np.empty((1, 2, 3, 3), dtype=np.float32))] rt_of(ins, outs) assert len(outs) == 1 expected_outpt = np.array([[[[10., 10., 10.], [10., 10., 10.], [10., 10., 10.]], [[26., 26., 26.], [26., 26., 26.], [26., 26., 26.]]]]) np.testing.assert_array_equal(outs[0], expected_outpt)
def test_vector_xbuffer_arg_name(self): def py_func(xbuffers): assert len(xbuffers) == 2 np.testing.assert_equal(xbuffers[0].to_numpy(), np.array([1., 2.], dtype=np.float32)) xbuffers[0].copy_from(xbuffers[0] * 2) np.testing.assert_equal(xbuffers[1].to_numpy(), np.array([-1., 0.], dtype=np.float32)) xbuffers[1].copy_from(xbuffers[1] * 2) xb1 = XBuffer(np.array([1., 2.], dtype=np.float32)) xb2 = XBuffer(np.array([-1., 0.], dtype=np.float32)) of = OpaqueFunc(py_func, [TypeCode.vXBuffer]) of([xb1, xb2]) np.testing.assert_equal(xb1.to_numpy(), np.array([2., 4.], dtype=np.float32)) np.testing.assert_equal(xb2.to_numpy(), np.array([-2., 0.], dtype=np.float32))
def test_xbuffer_construction(self): a = np.array([1., 1.5], dtype=np.float32) xb = XBuffer(a) np.testing.assert_equal(xb.to_numpy(), a)
class OpaqueFunc(object): # TypeCode conversion functions # First: C++ -> Python # Second: Python -> C++ type_codes_ = { TypeCode.vInt: ( lambda arg_: IntVector(arg_.ints), lambda arg_: lpx.OpaqueValue(lpx.IntVector(arg_))), TypeCode.Str: ( lambda arg_: arg_.s, lambda arg_: lpx.OpaqueValue(arg_)), TypeCode.Byte: ( lambda arg_: arg_.bytes, lambda arg_: lpx.OpaqueValue(arg_)), TypeCode.vStr: ( lambda arg_: StrVector(arg_.strings), lambda arg_: lpx.OpaqueValue(lpx.StrVector(arg_))), TypeCode.StrContainer: ( lambda arg_: StrContainer.from_lib(arg_.str_c), lambda arg_: lpx.OpaqueValue(arg_._str_c)), TypeCode.BytesContainer: ( lambda arg_: BytesContainer.from_lib(arg_.bytes_c), lambda arg_: lpx.OpaqueValue(arg_._bytes_c)), TypeCode.XGraph: ( lambda arg_: XGraph._from_xgraph(arg_.xg), lambda arg_: lpx.OpaqueValue(arg_._xgraph)), TypeCode.XBuffer: ( lambda arg_: XBuffer.from_lib(arg_.xb), lambda arg_: lpx.OpaqueValue(arg_._xb)), TypeCode.vXBuffer: ( lambda arg_: [XBuffer.from_lib(e) for e in arg_.xbuffers], lambda arg_: lpx.OpaqueValue( lpx.XBufferHolderVector([xb._xb for xb in arg_]))), TypeCode.OpaqueFunc: ( lambda arg_: OpaqueFunc.from_lib(arg_.of), lambda arg_: lpx.OpaqueValue(arg_._of)) } def __init__(self, func: Callable = None, type_codes: List[TypeCode] = None) -> None: self._of = lpx.OpaqueFunc() if type_codes is None: type_codes = [] if func is not None: self.set_func(func, type_codes) @classmethod def from_lib(cls, _of: lpx.OpaqueFunc) -> 'OpaqueFunc': of = OpaqueFunc.__new__(cls) of._of = _of return of def set_func(self, func: Callable, type_codes: List[TypeCode]): # if type_codes is not None: for tc in type_codes: if tc not in OpaqueFunc.type_codes_: raise NotImplementedError("Function with argument of" " unsupported type: {} provided" .format(tc.name)) def opaque_func_wrapper(args): new_args = [] if type_codes is not None: args_type_codes = type_codes else: args_type_codes = [TypeCode(args[i].get_type_code_int()) for i in range(len(args))] for tc, arg_ in zip(args_type_codes, args): if tc not in OpaqueFunc.type_codes_: raise ValueError(f"Unsupported type code: {tc}") arg_ = OpaqueFunc.type_codes_[tc][0](arg_) new_args.append(arg_) func(*new_args) arg_type_codes_ = lpx.IntVector([tc.value for tc in type_codes]) self._of.set_func(opaque_func_wrapper, arg_type_codes_) def __call__(self, *args: Any) -> None: """ Call internal lib OpaqueFunc with provided args """ args_type_codes = self.get_arg_type_codes() if len(args) != len(args_type_codes): raise ValueError("Invalid number of arguments detected." " OpaqueFunc is expecting {} arguments" " but got: {}" .format(len(args_type_codes), len(args))) oa_v = [] for tc, arg_ in zip(args_type_codes, args): if tc not in OpaqueFunc.type_codes_: raise ValueError(f"Unsupported type code: {tc}") oa_v.append(OpaqueFunc.type_codes_[tc][1](arg_)) oa = lpx.OpaqueArgs(oa_v) self._of(oa) def get_arg_type_codes(self): return [TypeCode(i) for i in self._of.get_arg_type_codes()] def get_nb_type_codes(self): return len(self.get_arg_type_codes()) def __del__(self): pass