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})
