text: # Lint as: python3
  # Copyright 2019 The IREE Authors
  #
  # Licensed under the Apache License v2.0 with LLVM Exceptions.
  # See https://llvm.org/LICENSE.txt for license information.
  # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
  
  import json
  import numpy as np
  
  from absl.testing import absltest
  
  from iree import runtime as rt
  from iree.runtime.function import (
      FunctionInvoker,
      IMPLICIT_BUFFER_ARG_MEMORY_TYPE,
      IMPLICIT_BUFFER_ARG_USAGE,
  )
  from iree.runtime.binding import VmVariantList
  
  
  class MockVmContext:
  
    def __init__(self, invoke_callback):
      self._invoke_callback = invoke_callback
      self.invocations = []
  
    def invoke(self, vm_function, arg_list, ret_list):
      self._invoke_callback(arg_list, ret_list)
      self.invocations.append((vm_function, arg_list, ret_list))
      print(f"INVOKE: {arg_list} -> {ret_list}")
  
    @property
    def mock_arg_reprs(self):
      return repr([arg_list for _, arg_list, _ in self.invocations])
  
  
  class MockVmFunction:
  
    def __init__(self, reflection):
      self.reflection = reflection
  
  
  class FunctionTest(absltest.TestCase):
  
    @classmethod
    def setUpClass(cls):
      # Doesn't matter what device. We just need one.
      config = rt.Config("vmvx")
      cls.device = config.device
  
    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 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 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 testListArgNoReflection(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([2, 3])
      self.assertEqual("[<VmVariantList(1): [List[2, 3]]>]",
                       vm_context.mock_arg_reprs)
  
    def testListArgArityMismatch(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)
      with self.assertRaisesRegex(ValueError,
                                  "expected a sequence with 2 values. got:"):
        _ = invoker([2, 3, 4])
  
    def testTupleArg(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": [["stuple", "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 testDictArg(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)
      _ = invoker({"b": 3, "a": 2})
      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})
  
    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 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 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 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 testMissingPositional(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.assertRaisesRegex(ValueError, "mismatched call arity:"):
        result = invoker(a=1, b=1)
  
    def testMissingPositionalNdarray(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(a=1, b=1)
  
    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.assertRaisesRegex(ValueError, "mismatched call arity:"):
        result = invoker(-1, a=1)
  
    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 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.assertRaisesRegex(ValueError, "specified kwarg 'c' is unknown"):
        result = invoker(-1, a=1, b=2, c=3)
  
    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 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 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 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 testDeviceArrayArgNoReflection(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={})
      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 testBufferViewArgNoReflection(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={})
      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 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 testReturnBufferViewNoReflection(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={})
      invoker = FunctionInvoker(vm_context, self.device, vm_function, tracer=None)
      result = invoker()
      np.testing.assert_array_equal([1, 0], result)
  
    # TODO: Fill out all return types.
    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 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))
  
  
  if __name__ == "__main__":
    absltest.main()
  