# 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

from typing import List
import iree.runtime
import iree.compiler

import gc
import numpy as np
import threading
import time
import unittest


class NonDeviceHalTest(unittest.TestCase):
    def testMemoryEnums(self):
        print("MemoryType:", iree.runtime.MemoryType)
        print("HOST_VISIBLE:", int(iree.runtime.MemoryType.HOST_VISIBLE))

        # Enum and/or operations on BufferCompatibility.
        self.assertEqual(
            iree.runtime.BufferCompatibility.IMPORTABLE
            | iree.runtime.BufferCompatibility.EXPORTABLE,
            int(iree.runtime.BufferCompatibility.IMPORTABLE)
            | int(iree.runtime.BufferCompatibility.EXPORTABLE),
        )
        self.assertEqual(
            iree.runtime.BufferCompatibility.EXPORTABLE
            & iree.runtime.BufferCompatibility.EXPORTABLE,
            int(iree.runtime.BufferCompatibility.EXPORTABLE),
        )

        # Enum and/or operations on BufferUsage.
        self.assertEqual(
            iree.runtime.BufferUsage.TRANSFER | iree.runtime.BufferUsage.MAPPING,
            int(iree.runtime.BufferUsage.TRANSFER)
            | int(iree.runtime.BufferUsage.MAPPING),
        )
        self.assertEqual(
            iree.runtime.BufferUsage.TRANSFER & iree.runtime.BufferUsage.TRANSFER,
            int(iree.runtime.BufferUsage.TRANSFER),
        )

        # Enum and/or operations on MemoryAccess.
        self.assertEqual(
            iree.runtime.MemoryAccess.READ | iree.runtime.MemoryAccess.WRITE,
            int(iree.runtime.MemoryAccess.READ) | int(iree.runtime.MemoryAccess.WRITE),
        )
        self.assertEqual(
            iree.runtime.MemoryAccess.ALL & iree.runtime.MemoryAccess.READ,
            int(iree.runtime.MemoryAccess.READ),
        )

        # Enum and/or operations on MemoryType.
        self.assertEqual(
            iree.runtime.MemoryType.DEVICE_LOCAL | iree.runtime.MemoryType.HOST_VISIBLE,
            int(iree.runtime.MemoryType.DEVICE_LOCAL)
            | int(iree.runtime.MemoryType.HOST_VISIBLE),
        )
        self.assertEqual(
            iree.runtime.MemoryType.OPTIMAL & iree.runtime.MemoryType.OPTIMAL,
            int(iree.runtime.MemoryType.OPTIMAL),
        )

    def testElementTypeEnums(self):
        i8 = iree.runtime.HalElementType.INT_8
        i4 = iree.runtime.HalElementType.INT_4
        self.assertTrue(iree.runtime.HalElementType.is_byte_aligned(i8))
        self.assertFalse(iree.runtime.HalElementType.is_byte_aligned(i4))
        self.assertEqual(1, iree.runtime.HalElementType.dense_byte_count(i8))


class DeviceHalTest(unittest.TestCase):
    def setUp(self):
        super().setUp()
        self.device = iree.runtime.get_device("local-task")
        self.allocator = self.device.allocator
        gc.collect()

    def testTrim(self):
        self.allocator.trim()
        # Just running is sufficient.

    def testProfilingDefaults(self):
        self.device.begin_profiling()
        self.device.flush_profiling()
        self.device.end_profiling()
        # Just running is sufficient.

    def testProfilingOptions(self):
        self.device.begin_profiling(mode="queue", file_path="foo.rdc")
        self.device.end_profiling()
        # Just running is sufficient.

    def testProfilingInvalidOptions(self):
        with self.assertRaisesRegex(ValueError, "unrecognized profiling mode"):
            self.device.begin_profiling(mode="SOMETHING THAT DOESN'T EXIST")

    def testStatistics(self):
        stats_dict = self.allocator.statistics
        stats_str = self.allocator.formatted_statistics
        if self.allocator.has_statistics:
            self.assertIn("host_bytes_peak", stats_dict)
            self.assertIn("host_bytes_allocated", stats_dict)
            self.assertIn("host_bytes_freed", stats_dict)
            self.assertIn("device_bytes_peak", stats_dict)
            self.assertIn("device_bytes_allocated", stats_dict)
            self.assertIn("device_bytes_freed", stats_dict)
            self.assertIn("HOST_LOCAL", stats_str)

    def testQueryCompatibility(self):
        compat = self.allocator.query_buffer_compatibility(
            memory_type=iree.runtime.MemoryType.DEVICE_LOCAL,
            allowed_usage=iree.runtime.BufferUsage.DEFAULT,
            intended_usage=iree.runtime.BufferUsage.DEFAULT,
            allocation_size=1024,
        )
        print("COMPAT:", compat)
        self.assertTrue(
            bool(compat & int(iree.runtime.BufferCompatibility.ALLOCATABLE)),
            "should be allocatable",
        )
        self.assertTrue(
            bool(compat & int(iree.runtime.BufferCompatibility.IMPORTABLE)),
            "should be importable",
        )
        self.assertTrue(
            bool(compat & int(iree.runtime.BufferCompatibility.EXPORTABLE)),
            "should be exportable",
        )

    def testAllocateBuffer(self):
        buffer = self.allocator.allocate_buffer(
            memory_type=iree.runtime.MemoryType.DEVICE_LOCAL,
            allowed_usage=iree.runtime.BufferUsage.DEFAULT,
            allocation_size=13,
        )
        print("BUFFER:", buffer)

    def testBufferViewConstructor(self):
        buffer = self.allocator.allocate_buffer(
            memory_type=iree.runtime.MemoryType.DEVICE_LOCAL,
            allowed_usage=iree.runtime.BufferUsage.DEFAULT,
            allocation_size=13,
        )
        bv = iree.runtime.HalBufferView(
            buffer, (1, 2), iree.runtime.HalElementType.INT_16
        )
        # NOTE: the exact bits set on type/usage/etc is implementation defined.
        self.assertEqual(
            repr(bv),
            "<HalBufferView (1, 2), element_type=0x10000010, 13 bytes (at offset 0 into 13), memory_type=DEVICE_LOCAL|HOST_VISIBLE, allowed_access=ALL, allowed_usage=TRANSFER|DISPATCH_STORAGE|MAPPING|MAPPING_PERSISTENT>",
        )
        self.assertEqual(4, bv.byte_length)

    def testBufferMap(self):
        buffer = self.allocator.allocate_buffer(
            memory_type=iree.runtime.MemoryType.DEVICE_LOCAL,
            allowed_usage=iree.runtime.BufferUsage.DEFAULT,
            allocation_size=13,
        )
        m = buffer.map()
        self.assertIsInstance(m, iree.runtime.MappedMemory)

    def testAllocateBufferCopy(self):
        ary = np.zeros([3, 4], dtype=np.int32) + 2
        buffer = self.allocator.allocate_buffer_copy(
            memory_type=iree.runtime.MemoryType.DEVICE_LOCAL,
            allowed_usage=iree.runtime.BufferUsage.DEFAULT,
            device=self.device,
            buffer=ary,
        )
        # NOTE: the exact bits set on type/usage/etc is implementation defined.
        self.assertEqual(
            repr(buffer),
            "<HalBuffer 48 bytes (at offset 0 into 48), memory_type=DEVICE_LOCAL|HOST_VISIBLE, allowed_access=ALL, allowed_usage=TRANSFER|DISPATCH_STORAGE|MAPPING|MAPPING_PERSISTENT>",
        )

    def testAllocateBufferViewCopy(self):
        ary = np.zeros([3, 4], dtype=np.int32) + 2
        buffer = self.allocator.allocate_buffer_copy(
            memory_type=iree.runtime.MemoryType.DEVICE_LOCAL,
            allowed_usage=iree.runtime.BufferUsage.DEFAULT,
            device=self.device,
            buffer=ary,
            element_type=iree.runtime.HalElementType.SINT_32,
        )
        # NOTE: the exact bits set on type/usage/etc is implementation defined.
        self.assertEqual(
            repr(buffer),
            "<HalBufferView (3, 4), element_type=0x20000011, 48 bytes (at offset 0 into 48), memory_type=DEVICE_LOCAL|HOST_VISIBLE, allowed_access=ALL, allowed_usage=TRANSFER|DISPATCH_STORAGE|MAPPING|MAPPING_PERSISTENT>",
        )

    def testAllocateHostStagingBufferCopy(self):
        buffer = self.allocator.allocate_host_staging_buffer_copy(
            self.device, np.int32(0)
        )
        # NOTE: the exact bits set on type/usage/etc is implementation defined.
        self.assertEqual(
            repr(buffer),
            "<HalBuffer 4 bytes (at offset 0 into 4), memory_type=DEVICE_LOCAL|HOST_VISIBLE, allowed_access=ALL, allowed_usage=TRANSFER|MAPPING|MAPPING_PERSISTENT>",
        )

    def testSemaphore(self):
        sem0 = self.device.create_semaphore(0)
        self.assertEqual(sem0.query(), 0)
        sem1 = self.device.create_semaphore(1)
        self.assertEqual(sem1.query(), 1)
        sem1.signal(2)
        self.assertEqual(sem1.query(), 2)

    def testSemaphoreSignal(self):
        sem = self.device.create_semaphore(0)
        self.assertFalse(sem.wait(1, deadline=0))
        sem.signal(1)
        self.assertTrue(sem.wait(1, deadline=0))

    def testSynchronousSemaphoreFailed(self):
        sem = self.device.create_semaphore(0)
        sem.fail("TEST FAILURE")
        with self.assertRaisesRegex(
            RuntimeError, "^synchronous semaphore failure.*TEST FAILURE"
        ):
            sem.wait(1, deadline=0)

    def testAsynchronousSemaphoreFailed(self):
        sem = self.device.create_semaphore(0)
        exceptions = []

        def run():
            print("SIGNALLING ASYNC FAILURE")
            time.sleep(0.2)
            sem.fail("TEST FAILURE")
            print("SIGNALLED")

        def wait():
            print("WAITING")
            try:
                sem.wait(1)
            except RuntimeError as e:
                exceptions.append(e)

        runner = threading.Thread(target=run)
        waiter = threading.Thread(target=wait)
        waiter.start()
        runner.start()
        waiter.join()
        runner.join()
        self.assertTrue(exceptions)
        print(exceptions)
        # Note: It is impossible to 100% guarantee that this sequences such as to
        # report an asynchronous vs synchronous failure, although we tip the odds in
        # this favor with the sleep in the signalling thread. Therefore, we do not
        # check the "asynchronous" vs "synchronous" message prefix to avoid flaky
        # test races.
        self.assertIn("TEST FAILURE", str(exceptions[0]))

    def testTrivialQueueAlloc(self):
        sem = self.device.create_semaphore(0)
        buf = self.device.queue_alloca(
            1024, wait_semaphores=[(sem, 0)], signal_semaphores=[(sem, 1)]
        )
        self.assertIsInstance(buf, iree.runtime.HalBuffer)
        self.device.queue_dealloca(
            buf, wait_semaphores=[(sem, 1)], signal_semaphores=[]
        )

    def testAllocAcceptsFences(self):
        # Also tests HalFence, HalFence.insert, HalFence.wait (infinite)
        sem = self.device.create_semaphore(0)
        fence0 = iree.runtime.HalFence(1)
        fence0.insert(sem, 0)
        fence1 = iree.runtime.HalFence(1)
        fence1.insert(sem, 1)
        fence2 = iree.runtime.HalFence(2)
        fence2.insert(sem, 2)
        buf = self.device.queue_alloca(
            1024, wait_semaphores=fence0, signal_semaphores=fence1
        )
        self.assertIsInstance(buf, iree.runtime.HalBuffer)
        self.device.queue_dealloca(
            buf, wait_semaphores=fence1, signal_semaphores=fence2
        )
        self.assertTrue(fence2.wait())
        self.assertEqual(sem.query(), 2)

    def testFenceCreateAt(self):
        sem = self.device.create_semaphore(0)
        fence = iree.runtime.HalFence.create_at(sem, 1)
        self.assertFalse(fence.wait(deadline=0))
        sem.signal(1)
        self.assertTrue(fence.wait(deadline=0))

    def testFenceSignal(self):
        sem = self.device.create_semaphore(0)
        fence = iree.runtime.HalFence.create_at(sem, 1)
        self.assertFalse(fence.wait(deadline=0))
        fence.signal()
        self.assertTrue(fence.wait(deadline=0))

    def testSynchronousFenceFailed(self):
        sem = self.device.create_semaphore(0)
        fence = iree.runtime.HalFence.create_at(sem, 1)
        fence.fail("TEST FAILURE")
        with self.assertRaisesRegex(
            RuntimeError, "^synchronous fence failure.*TEST FAILURE"
        ):
            fence.wait(deadline=0)

    def testAsynchronousFenceFailed(self):
        sem = self.device.create_semaphore(0)
        fence = iree.runtime.HalFence.create_at(sem, 1)
        exceptions = []

        def run():
            print("SIGNALLING ASYNC FAILURE")
            time.sleep(0.2)
            fence.fail("TEST FAILURE")
            print("SIGNALLED")

        def wait():
            print("WAITING")
            try:
                fence.wait()
            except RuntimeError as e:
                exceptions.append(e)

        runner = threading.Thread(target=run)
        waiter = threading.Thread(target=wait)
        waiter.start()
        runner.start()
        waiter.join()
        runner.join()
        self.assertTrue(exceptions)
        print(exceptions)
        # Note: It is impossible to 100% guarantee that this sequences such as to
        # report an asynchronous vs synchronous failure, although we tip the odds in
        # this favor with the sleep in the signalling thread. Therefore, we do not
        # check the "asynchronous" vs "synchronous" message prefix to avoid flaky
        # test races.
        self.assertIn("TEST FAILURE", str(exceptions[0]))

    def testFenceJoin(self):
        sem1 = self.device.create_semaphore(0)
        sem2 = self.device.create_semaphore(0)
        fence1 = iree.runtime.HalFence.create_at(sem1, 1)
        fence2 = iree.runtime.HalFence.create_at(sem2, 1)
        fence = iree.runtime.HalFence.join([fence1, fence2])
        self.assertEqual(fence.timepoint_count, 2)

    def testFenceInsert(self):
        sem1 = self.device.create_semaphore(0)
        sem2 = self.device.create_semaphore(0)
        fence = iree.runtime.HalFence(2)
        fence.insert(sem1, 1)
        self.assertEqual(fence.timepoint_count, 1)
        fence.insert(sem1, 2)
        self.assertEqual(fence.timepoint_count, 1)
        fence.insert(sem2, 2)
        self.assertEqual(fence.timepoint_count, 2)

    def testFenceExtend(self):
        sem1 = self.device.create_semaphore(0)
        sem2 = self.device.create_semaphore(0)
        fence = iree.runtime.HalFence(2)
        fence.insert(sem1, 1)
        self.assertEqual(fence.timepoint_count, 1)
        fence.extend(iree.runtime.HalFence.create_at(sem2, 2))
        self.assertEqual(fence.timepoint_count, 2)

    def testRoundTripQueueCopy(self):
        original_ary = np.zeros([3, 4], dtype=np.int32) + 2
        source_bv = self.allocator.allocate_buffer_copy(
            memory_type=iree.runtime.MemoryType.DEVICE_LOCAL,
            allowed_usage=iree.runtime.BufferUsage.DEFAULT,
            device=self.device,
            buffer=original_ary,
            element_type=iree.runtime.HalElementType.SINT_32,
        )
        source_buffer = source_bv.get_buffer()
        target_buffer = self.allocator.allocate_buffer(
            memory_type=iree.runtime.MemoryType.DEVICE_LOCAL,
            allowed_usage=iree.runtime.BufferUsage.DEFAULT,
            allocation_size=source_buffer.byte_length(),
        )
        sem = self.device.create_semaphore(0)
        self.device.queue_copy(
            source_buffer,
            target_buffer,
            wait_semaphores=iree.runtime.HalFence.create_at(sem, 0),
            signal_semaphores=iree.runtime.HalFence.create_at(sem, 1),
        )
        iree.runtime.HalFence.create_at(sem, 1).wait()
        copy_ary = target_buffer.map().asarray(original_ary.shape, original_ary.dtype)
        np.testing.assert_array_equal(original_ary, copy_ary)

    def testIncompatibleSizeQueueCopy(self):
        source_buffer = self.allocator.allocate_buffer(
            memory_type=iree.runtime.MemoryType.DEVICE_LOCAL,
            allowed_usage=iree.runtime.BufferUsage.DEFAULT,
            allocation_size=13,
        )
        target_buffer = self.allocator.allocate_buffer(
            memory_type=iree.runtime.MemoryType.DEVICE_LOCAL,
            allowed_usage=iree.runtime.BufferUsage.DEFAULT,
            allocation_size=12,
        )
        sem = self.device.create_semaphore(0)
        with self.assertRaisesRegex(ValueError, "length must be less than"):
            self.device.queue_copy(
                source_buffer,
                target_buffer,
                wait_semaphores=iree.runtime.HalFence.create_at(sem, 0),
                signal_semaphores=iree.runtime.HalFence.create_at(sem, 1),
            )

    def testCommandBufferStartsByDefault(self):
        cb = iree.runtime.HalCommandBuffer(self.device)
        with self.assertRaisesRegex(RuntimeError, "FAILED_PRECONDITION"):
            cb.begin()
        cb = iree.runtime.HalCommandBuffer(self.device, begin=False)
        cb.begin()

    def testCommandBufferCopy(self):
        # Doesn't test much but that calls succeed.
        cb = iree.runtime.HalCommandBuffer(self.device)
        buffer1 = self.allocator.allocate_buffer(
            memory_type=iree.runtime.MemoryType.DEVICE_LOCAL,
            allowed_usage=iree.runtime.BufferUsage.DEFAULT,
            allocation_size=13,
        )
        buffer2 = self.allocator.allocate_buffer(
            memory_type=iree.runtime.MemoryType.DEVICE_LOCAL,
            allowed_usage=iree.runtime.BufferUsage.DEFAULT,
            allocation_size=13,
        )
        cb.copy(buffer1, buffer2, end=True)
        with self.assertRaisesRegex(RuntimeError, "FAILED_PRECONDITION"):
            cb.end()

    def testCommandBufferFill(self):
        # Doesn't test much but that calls succeed.
        cb = iree.runtime.HalCommandBuffer(self.device)
        buffer1 = self.allocator.allocate_buffer(
            memory_type=iree.runtime.MemoryType.DEVICE_LOCAL,
            allowed_usage=iree.runtime.BufferUsage.DEFAULT,
            allocation_size=12,
        )
        cb.fill(buffer1, np.int32(1), 0, 12, end=True)
        with self.assertRaisesRegex(RuntimeError, "FAILED_PRECONDITION"):
            cb.end()

    def testCommandBufferExecute(self):
        # Doesn't test much but that calls succeed.
        cb = iree.runtime.HalCommandBuffer(self.device)
        buffer1 = self.allocator.allocate_buffer(
            memory_type=iree.runtime.MemoryType.DEVICE_LOCAL,
            allowed_usage=iree.runtime.BufferUsage.DEFAULT,
            allocation_size=12,
        )
        cb.fill(buffer1, np.int32(1), 0, 12, end=True)

        sem = self.device.create_semaphore(0)
        self.device.queue_execute(
            cb, wait_semaphores=[(sem, 0)], signal_semaphores=[(sem, 1)]
        )
        iree.runtime.HalFence.create_at(sem, 1).wait()

    def testCommandBufferExecuteAcceptsFence(self):
        # Doesn't test much but that calls succeed.
        cb = iree.runtime.HalCommandBuffer(self.device)
        buffer1 = self.allocator.allocate_buffer(
            memory_type=iree.runtime.MemoryType.DEVICE_LOCAL,
            allowed_usage=iree.runtime.BufferUsage.DEFAULT,
            allocation_size=12,
        )
        cb.fill(buffer1, np.int32(1), 0, 12, end=True)

        sem = self.device.create_semaphore(0)
        self.device.queue_execute(
            cb,
            wait_semaphores=iree.runtime.HalFence.create_at(sem, 0),
            signal_semaphores=iree.runtime.HalFence.create_at(sem, 1),
        )
        iree.runtime.HalFence.create_at(sem, 1).wait()


class DeviceDLPackTest(unittest.TestCase):
    """Tests low level DLPack import/export against the CPU HAL backend.

    This test leverages the fact that numpy is a reasonable dlpack
    producer/consumer. It has the caveat that our low level support does not
    allow import of non page aligned data, so we have to take some extra
    steps to prep it. For pure CPU/Numpy import/export, we have better
    supported paths than this, but we leverage it here for its testing
    value, as it exercises code paths that are otherwise only accessible
    on devices.
    """

    def setUp(self):
        super().setUp()
        self.device = iree.runtime.get_device("local-task")
        self.allocator = self.device.allocator
        gc.collect()

    def roundtrip(self, input_array, element_type):
        # We have top copy the input array into our own buffer to ensure
        # alignment (dlpack import/export require aligned data).
        orig_bv = self.allocator.allocate_buffer_copy(
            memory_type=iree.runtime.MemoryType.DEVICE_LOCAL,
            allowed_usage=iree.runtime.BufferUsage.DEFAULT,
            device=self.device,
            buffer=input_array,
            element_type=element_type,
        )
        aligned_input_array = orig_bv.map().asarray(
            input_array.shape, input_array.dtype
        )

        # Export the __dlpack__ capsule from numpy, which should be a plain
        # view over the buffer we originally allocated (therefore, aligned
        # and importable).
        input_capsule = aligned_input_array.__dlpack__()
        aligned_input_array = None
        gc.collect()
        imported_bv = self.device.from_dlpack_capsule(input_capsule)

        # Export a capsule from this imported buffer view and create a new
        # array out of it.
        class DummyProducer:
            def __dlpack__(_, stream=None):
                capsule = self.device.create_dlpack_capsule(imported_bv, 1, 0)
                return capsule

            def __dlpack_device__(self):
                return (1, 0)  # CPU, id 0

        reimported_array = np.from_dlpack(DummyProducer())
        imported_bv = None
        gc.collect()
        np.testing.assert_array_equal(input_array, reimported_array)

    def testImportExportF64(self):
        self.roundtrip(np.random.rand(3, 4), iree.runtime.HalElementType.FLOAT_64)

    def testImportExportF32(self):
        self.roundtrip(
            np.random.rand(3, 4, 16, 32, 1, 5, 2).astype(np.float32),
            iree.runtime.HalElementType.FLOAT_32,
        )

    def testImportExportF16(self):
        self.roundtrip(
            np.random.rand(3, 4).astype(np.float16),
            iree.runtime.HalElementType.FLOAT_16,
        )

    def testImportExportSI8(self):
        self.roundtrip(
            (np.random.rand(3, 4) * 255.0).astype(np.int8),
            iree.runtime.HalElementType.SINT_8,
        )

    def testImportExportSI16(self):
        self.roundtrip(
            (np.random.rand(3, 4) * 255.0).astype(np.int16),
            iree.runtime.HalElementType.SINT_16,
        )

    def testImportExportSI32(self):
        self.roundtrip(
            (np.random.rand(3, 4) * 255.0).astype(np.int32),
            iree.runtime.HalElementType.SINT_32,
        )

    def testImportExportSI64(self):
        self.roundtrip(
            (np.random.rand(3, 4) * 255.0).astype(np.int64),
            iree.runtime.HalElementType.SINT_64,
        )

    def testImportExportUI8(self):
        self.roundtrip(
            (np.random.rand(3, 4) * 255.0).astype(np.uint8),
            iree.runtime.HalElementType.UINT_8,
        )

    def testImportExportUI16(self):
        self.roundtrip(
            (np.random.rand(3, 4) * 255.0).astype(np.uint16),
            iree.runtime.HalElementType.UINT_16,
        )

    def testImportExportUI32(self):
        self.roundtrip(
            (np.random.rand(3, 4) * 255.0).astype(np.uint32),
            iree.runtime.HalElementType.UINT_32,
        )

    def testImportExportUI64(self):
        self.roundtrip(
            (np.random.rand(3, 4) * 255.0).astype(np.uint64),
            iree.runtime.HalElementType.UINT_64,
        )

    def testImportExportBool(self):
        self.roundtrip(
            (np.random.rand(3, 4) * 255.0).astype(np.bool_),
            iree.runtime.HalElementType.BOOL_8,
        )

    def testImportExportUI64(self):
        self.roundtrip(
            (np.random.rand(3, 4) * 255.0).astype(np.uint64),
            iree.runtime.HalElementType.UINT_64,
        )

    def testImportExportComplex64(self):
        shape = (3, 1, 5, 6, 12, 2, 3)
        self.roundtrip(
            np.random.uniform(-1, 1, shape) + 1.0j * np.random.uniform(-1, 1, shape),
            iree.runtime.HalElementType.COMPLEX_64,
        )

    def testImportExportComplex64(self):
        shape = (3, 1, 5, 6, 12, 2, 3)
        self.roundtrip(
            (
                np.random.uniform(-1, 1, shape) + 1.0j * np.random.uniform(-1, 1, shape)
            ).astype(np.complex128),
            iree.runtime.HalElementType.COMPLEX_64,
        )


class HalModuleDebugSinkTest(unittest.TestCase):
    COMPILED_TRACE_TENSOR: bytes

    @classmethod
    def compile_trace_tensor(cls):
        if not hasattr(cls, "COMPILED_TRACE_TENSOR"):
            cls.COMPILED_TRACE_TENSOR = iree.compiler.compile_str(
                """
                func.func @trace_args(%arg0: tensor<2xi32>, %arg1: tensor<3xi32>) {
                    flow.tensor.trace "debug_sink_test" = [
                        %arg0: tensor<2xi32>,
                        %arg1: tensor<3xi32>
                    ]
                    return
                }
                """,
                target_backends=iree.compiler.core.DEFAULT_TESTING_BACKENDS,
            )
        return cls.COMPILED_TRACE_TENSOR

    def testHalModuleBufferViewTraceCallback(self):
        """Check that the trace tensor callback gets called with the expected
        arguments."""
        program_bytes = HalModuleDebugSinkTest.compile_trace_tensor()

        arg0 = np.array([1, 2], dtype=np.int32)
        arg1 = np.array([3, 4, 5], dtype=np.int32)

        callback_key: str = None
        callback_buffer_views = None

        def callback(key: str, buffer_views: List[iree.runtime.HalBufferView]):
            nonlocal callback_key
            callback_key = key
            nonlocal callback_buffer_views
            callback_buffer_views = buffer_views

        instance = iree.runtime.VmInstance()
        device = iree.runtime.get_device(iree.compiler.core.DEFAULT_TESTING_DRIVER)
        hal_module = iree.runtime.create_hal_module(
            instance, device, debug_sink=iree.runtime.HalModuleDebugSink(callback)
        )
        program_module = iree.runtime.VmModule.copy_buffer(instance, program_bytes)
        context = iree.runtime.VmContext(instance)
        context.register_modules([hal_module, program_module])
        fn = program_module.lookup_function("trace_args")
        fn_invoker = iree.runtime.FunctionInvoker(context, device, fn)
        fn_invoker(arg0, arg1)

        assert callback_key == "debug_sink_test"
        assert len(callback_buffer_views) == 2
        actual_arg0 = iree.runtime.DeviceArray(
            device, callback_buffer_views[0]
        ).to_host()
        actual_arg1 = iree.runtime.DeviceArray(
            device, callback_buffer_views[1]
        ).to_host()
        np.testing.assert_equal(actual_arg0, arg0)
        np.testing.assert_equal(actual_arg1, arg1)

    def testNoneHalModuleDebugSink(self):
        device = iree.runtime.get_device(iree.compiler.core.DEFAULT_TESTING_DRIVER)
        instance = iree.runtime.VmInstance()
        hal_module = iree.runtime.create_hal_module(
            instance,
            device,
            debug_sink=None,
        )

    def testExceptionInHalModuleBufferViewTraceCallback(self):
        """When an exception occurs in the callback check that it properly propagates
        through the bindings and results in a IREE module function failed invocation.
        """
        program_bytes = HalModuleDebugSinkTest.compile_trace_tensor()

        arg0 = np.array([1, 2], dtype=np.int32)
        arg1 = np.array([3, 4, 5], dtype=np.int32)

        device = iree.runtime.get_device(iree.compiler.core.DEFAULT_TESTING_DRIVER)

        class TestException(Exception):
            def __init__(self, msg: str):
                super().__init__(msg)

        def callback(key: str, buffer_views: List[iree.runtime.HalBufferView]):
            raise TestException("This is a test exception")

        instance = iree.runtime.VmInstance()
        hal_module = iree.runtime.create_hal_module(
            instance, device, debug_sink=iree.runtime.HalModuleDebugSink(callback)
        )
        program_module = iree.runtime.VmModule.copy_buffer(instance, program_bytes)
        context = iree.runtime.VmContext(instance)
        context.register_modules([hal_module, program_module])
        fn = program_module.lookup_function("trace_args")
        fn_invoker = iree.runtime.FunctionInvoker(context, device, fn)
        # TODO: once IREE status chains messages test for the actual message we raise
        # within the callback.
        self.assertRaisesRegex(RuntimeError, "UNKNOWN", fn_invoker, arg0, arg1)

    def testHalModuleBufferViewTraceCallbackReferencingItselfDoesNotLeak(self):
        """Check that if we do not hold reference to the HAL module or VM context,
        but we hold a reference to the debug sink in the callback, the callback object
        does not leak.
        """
        is_callback_destroyed: bool = False

        class Callback:
            def __del__(self):
                nonlocal is_callback_destroyed
                is_callback_destroyed = True

            def __call__(
                self, key: str, buffer_views: List[iree.runtime.HalBufferView]
            ):
                pass

        callback = Callback()
        debug_sink = iree.runtime.HalModuleDebugSink(callback)
        setattr(callback, "debug_sink", debug_sink)

        device = iree.runtime.get_device(iree.compiler.core.DEFAULT_TESTING_DRIVER)

        vm_instance = iree.runtime.VmInstance()
        hal_module = iree.runtime.create_hal_module(
            vm_instance,
            device,
            debug_sink=debug_sink,
        )
        vm_context = iree.runtime.VmContext(vm_instance)
        vm_context.register_modules([hal_module])
        assert not is_callback_destroyed

        del callback
        del debug_sink
        del hal_module
        del vm_instance
        del vm_context
        gc.collect()
        assert is_callback_destroyed


if __name__ == "__main__":
    unittest.main()