def testNestedResults(self): def invoke(arg_list, ret_list): ret_list.push_int(3) sub_list = VmVariantList(2) sub_dict = VmVariantList(2) sub_dict.push_int(100) sub_dict.push_int(200) sub_list.push_list(sub_dict) sub_list.push_int(6) ret_list.push_list(sub_list) vm_context = MockVmContext(invoke) vm_function = MockVmFunction( reflection={ "iree.abi": json.dumps({ "a": [], "r": [ "i32", [ "slist", ["sdict", ["bar", "i32"], ["foo", "i32"]], "i64", ] ], }) }) invoker = FunctionInvoker(vm_context, self.device, vm_function, tracer=None) result = invoker() self.assertEqual((3, [{'bar': 100, 'foo': 200}, 6]), result)
def testReturnBufferView(self): result_array = np.asarray([1, 0], dtype=np.int32) def invoke(arg_list, ret_list): buffer_view = self.device.allocator.allocate_buffer_copy( memory_type=IMPLICIT_BUFFER_ARG_MEMORY_TYPE, allowed_usage=IMPLICIT_BUFFER_ARG_USAGE, buffer=result_array, element_type=rt.HalElementType.SINT_32) ret_list.push_buffer_view(buffer_view) vm_context = MockVmContext(invoke) vm_function = MockVmFunction(reflection={ "iree.abi": json.dumps({ "a": [], "r": [["ndarray", "i32", 1, 2]], }) }) invoker = FunctionInvoker(vm_context, self.device, vm_function, tracer=None) result = invoker() np.testing.assert_array_equal([1, 0], result)
def testKeywordArgs(self): def invoke(arg_list, ret_list): ret_list.push_int(3) vm_context = MockVmContext(invoke) vm_function = MockVmFunction( reflection={ "iree.abi": json.dumps({ "a": [ "i32", ["named", "a", "i32"], ["named", "b", "i32"], ], "r": [ "i32", ], }) }) invoker = FunctionInvoker(vm_context, self.device, vm_function, tracer=None) result = invoker(-1, a=1, b=2) self.assertEqual("[<VmVariantList(3): [-1, 1, 2]>]", vm_context.mock_arg_reprs) self.assertEqual(3, result)
def testBufferViewArg(self): arg_buffer_view = self.device.allocator.allocate_buffer_copy( memory_type=IMPLICIT_BUFFER_ARG_MEMORY_TYPE, allowed_usage=IMPLICIT_BUFFER_ARG_USAGE, buffer=np.asarray([1, 0], dtype=np.int32), element_type=rt.HalElementType.SINT_32) invoked_arg_list = None def invoke(arg_list, ret_list): nonlocal invoked_arg_list invoked_arg_list = arg_list vm_context = MockVmContext(invoke) vm_function = MockVmFunction(reflection={ "iree.abi": json.dumps({ "a": [["ndarray", "i32", 1, 2]], "r": [], }) }) invoker = FunctionInvoker(vm_context, self.device, vm_function, tracer=None) _ = invoker(arg_buffer_view) self.assertEqual("<VmVariantList(1): [HalBufferView(2:0x20000011)]>", repr(invoked_arg_list))
def testMissingKeyword(self): def invoke(arg_list, ret_list): ret_list.push_int(3) vm_context = MockVmContext(invoke) vm_function = MockVmFunction( reflection={ "iree.abi": json.dumps({ "a": [ "i32", ["named", "a", "i32"], ["named", "b", "i32"], ], "r": [ "i32", ], }) }) invoker = FunctionInvoker(vm_context, self.device, vm_function, tracer=None) with self.assertRaisesRegexp(ValueError, "a required argument was not specified"): result = invoker(-1, a=1)
def testDeviceArrayArg(self): # Note that since the device array is set up to disallow implicit host # transfers, this also verifies that no accidental/automatic transfers # are done as part of marshalling the array to the function. arg_array = rt.asdevicearray(self.device, np.asarray([1, 0], dtype=np.int32), implicit_host_transfer=False) invoked_arg_list = None def invoke(arg_list, ret_list): nonlocal invoked_arg_list invoked_arg_list = arg_list vm_context = MockVmContext(invoke) vm_function = MockVmFunction(reflection={ "iree.abi": json.dumps({ "a": [["ndarray", "i32", 1, 2]], "r": [], }) }) invoker = FunctionInvoker(vm_context, self.device, vm_function, tracer=None) result = invoker(arg_array) self.assertEqual("<VmVariantList(1): [HalBufferView(2:0x20000011)]>", repr(invoked_arg_list))
def testExtraKeyword(self): def invoke(arg_list, ret_list): ret_list.push_int(3) vm_context = MockVmContext(invoke) vm_function = MockVmFunction( reflection={ "iree.abi": json.dumps({ "a": [ "i32", ["named", "a", "i32"], ["named", "b", "i32"], ], "r": [ "i32", ], }) }) invoker = FunctionInvoker(vm_context, self.device, vm_function, tracer=None) with self.assertRaisesRegexp(ValueError, "specified kwarg 'c' is unknown"): result = invoker(-1, a=1, b=2, c=3)
def testMissingKeywordNdArray(self): def invoke(arg_list, ret_list): ret_list.push_int(3) vm_context = MockVmContext(invoke) vm_function = MockVmFunction( reflection={ "iree.abi": json.dumps({ "a": [ ["ndarray", "i32", 1, 1], ["named", "a", ["ndarray", "i32", 1, 1]], ["named", "b", ["ndarray", "i32", 1, 1]], ], "r": [ "i32", ], }) }) invoker = FunctionInvoker(vm_context, self.device, vm_function, tracer=None) with self.assertRaisesRegex(ValueError, "mismatched call arity:"): result = invoker(-1, a=1)
def testReturnTypeNdArrayBool(self): result_array = np.asarray([1, 0], dtype=np.int8) def invoke(arg_list, ret_list): buffer_view = self.device.allocator.allocate_buffer_copy( memory_type=IMPLICIT_BUFFER_ARG_MEMORY_TYPE, allowed_usage=IMPLICIT_BUFFER_ARG_USAGE, buffer=result_array, element_type=rt.HalElementType.UINT_8) ret_list.push_buffer_view(buffer_view) vm_context = MockVmContext(invoke) vm_function = MockVmFunction( reflection={ "iree.abi": json.dumps({ "a": [], "r": [["ndarray", "i1", 1, 2]], }) }) invoker = FunctionInvoker(vm_context, self.device, vm_function, tracer=None) result = invoker() # assertEqual on bool arrays is fraught for... reasons. np.testing.assert_array_equal([True, False], result)
def testDictArgNoReflection(self): def invoke(arg_list, ret_list): ret_list.push_int(3) vm_context = MockVmContext(invoke) vm_function = MockVmFunction(reflection={}) invoker = FunctionInvoker(vm_context, self.device, vm_function, tracer=None) _ = invoker({"b": 3, "a": 2}) self.assertEqual("[<VmVariantList(1): [List[2, 3]]>]", vm_context.mock_arg_reprs)
def testNoReflectionScalars(self): def invoke(arg_list, ret_list): ret_list.push_int(3) ret_list.push_int(4) vm_context = MockVmContext(invoke) vm_function = MockVmFunction(reflection={}) invoker = FunctionInvoker(vm_context, self.device, vm_function, tracer=None) result = invoker(1, 2) self.assertEqual("[<VmVariantList(2): [1, 2]>]", vm_context.mock_arg_reprs) self.assertEqual((3, 4), result)
def testNdarrayArgNoReflection(self): arg_array = np.asarray([1, 0], dtype=np.int32) invoked_arg_list = None def invoke(arg_list, ret_list): nonlocal invoked_arg_list invoked_arg_list = arg_list vm_context = MockVmContext(invoke) vm_function = MockVmFunction(reflection={}) invoker = FunctionInvoker(vm_context, self.device, vm_function, tracer=None) result = invoker(arg_array) self.assertEqual("<VmVariantList(1): [HalBufferView(2:0x20000011)]>", repr(invoked_arg_list))
def testInlinedResults(self): def invoke(arg_list, ret_list): ret_list.push_int(3) ret_list.push_int(4) vm_context = MockVmContext(invoke) vm_function = MockVmFunction( reflection={ "iree.abi": json.dumps({ "a": [], "r": [["slist", "i32", "i32"]], }) }) invoker = FunctionInvoker(vm_context, self.device, vm_function, tracer=None) result = invoker() self.assertEqual([3, 4], result)
def testReturnTypeNdArrayBool(self): result_array = np.asarray([1, 0], dtype=np.int8) def invoke(arg_list, ret_list): ret_list.push_buffer_view(self.device, result_array, rt.HalElementType.UINT_8) vm_context = MockVmContext(invoke) vm_function = MockVmFunction( reflection={ "iree.abi": json.dumps({ "a": [], "r": [["ndarray", "i1", 1, 2]], }) }) invoker = FunctionInvoker(vm_context, self.device, vm_function, tracer=None) result = invoker() # assertEqual on bool arrays is fraught for... reasons. self.assertEqual("array([ True, False])", repr(result))
def testReturnTypeList(self): vm_list = VmVariantList(2) vm_list.push_int(1) vm_list.push_int(2) def invoke(arg_list, ret_list): ret_list.push_list(vm_list) vm_context = MockVmContext(invoke) vm_function = MockVmFunction( reflection={ "iree.abi": json.dumps({ "a": [], "r": [["py_homogeneous_list", "i64"]], }) }) invoker = FunctionInvoker(vm_context, self.device, vm_function, tracer=None) result = invoker() self.assertEqual("[1, 2]", repr(result))
def testNdarrayArg(self): arg_array = np.asarray([1, 0], dtype=np.int32) invoked_arg_list = None def invoke(arg_list, ret_list): nonlocal invoked_arg_list invoked_arg_list = arg_list vm_context = MockVmContext(invoke) vm_function = MockVmFunction(reflection={ "iree.abi": json.dumps({ "a": [["ndarray", "i32", 1, 2]], "r": [], }) }) invoker = FunctionInvoker(vm_context, self.device, vm_function, tracer=None) result = invoker(arg_array) self.assertEqual("<VmVariantList(1): [HalBufferView(2:0x20000011)]>", repr(invoked_arg_list))
def testDictArgKeyError(self): def invoke(arg_list, ret_list): ret_list.push_int(3) vm_context = MockVmContext(invoke) vm_function = MockVmFunction( reflection={ "iree.abi": json.dumps({ "a": [ ["sdict", ["a", "i32"], ["b", "i32"]], ], "r": [ "i32", ], }) }) invoker = FunctionInvoker(vm_context, self.device, vm_function, tracer=None) with self.assertRaisesRegex(ValueError, "could not get item 'b' from: "): _ = invoker({"a": 2, "c": 3})
def testListArg(self): def invoke(arg_list, ret_list): ret_list.push_int(3) vm_context = MockVmContext(invoke) vm_function = MockVmFunction( reflection={ "iree.abi": json.dumps({ "a": [ ["slist", "i32", "i32"], ], "r": [ "i32", ], }) }) invoker = FunctionInvoker(vm_context, self.device, vm_function, tracer=None) _ = invoker([2, 3]) self.assertEqual("[<VmVariantList(1): [List[2, 3]]>]", vm_context.mock_arg_reprs)
def testDictArgArityMismatch(self): def invoke(arg_list, ret_list): ret_list.push_int(3) vm_context = MockVmContext(invoke) vm_function = MockVmFunction( reflection={ "iree.abi": json.dumps({ "a": [ ["sdict", ["a", "i32"], ["b", "i32"]], ], "r": [ "i32", ], }) }) invoker = FunctionInvoker(vm_context, self.device, vm_function, tracer=None) with self.assertRaisesRegex(ValueError, "expected a dict with 2 values. got:"): _ = invoker({"a": 2, "b": 3, "c": 4})