pytorch3d/tests/test_opengl_utils.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

import ctypes
import os
import sys
import threading
import unittest

import torch

os.environ["PYOPENGL_PLATFORM"] = "egl"
import pycuda._driver  # noqa
from OpenGL import GL as gl  # noqa
from OpenGL.raw.EGL._errors import EGLError  # noqa
from pytorch3d.renderer.opengl import _can_import_egl_and_pycuda  # noqa
from pytorch3d.renderer.opengl.opengl_utils import (  # noqa
    _define_egl_extension,
    _egl_convert_to_int_array,
    _get_cuda_device,
    egl,
    EGLContext,
    global_device_context_store,
)

from .common_testing import TestCaseMixin, usesOpengl  # noqa

MAX_EGL_HEIGHT = global_device_context_store.max_egl_height
MAX_EGL_WIDTH = global_device_context_store.max_egl_width


def _draw_square(r=1.0, g=0.0, b=1.0, **kwargs) -> torch.Tensor:
    gl.glClear(gl.GL_COLOR_BUFFER_BIT)
    gl.glColor3f(r, g, b)
    x1, x2 = -0.5, 0.5
    y1, y2 = -0.5, 0.5
    gl.glRectf(x1, y1, x2, y2)
    out_buffer = gl.glReadPixels(
        0, 0, MAX_EGL_WIDTH, MAX_EGL_HEIGHT, gl.GL_RGB, gl.GL_UNSIGNED_BYTE
    )
    image = torch.frombuffer(out_buffer, dtype=torch.uint8).reshape(
        MAX_EGL_HEIGHT, MAX_EGL_WIDTH, 3
    )
    return image


def _draw_squares_with_context(
    cuda_device_id=0, result=None, thread_id=None, **kwargs
) -> None:
    context = EGLContext(MAX_EGL_WIDTH, MAX_EGL_HEIGHT, cuda_device_id)
    with context.active_and_locked():
        images = []
        for _ in range(3):
            images.append(_draw_square(**kwargs).float())
        if result is not None and thread_id is not None:
            egl_info = context.get_context_info()
            data = {"egl": egl_info, "images": images}
            result[thread_id] = data


def _draw_squares_with_context_store(
    cuda_device_id=0,
    result=None,
    thread_id=None,
    verbose=False,
    **kwargs,
) -> None:
    device = torch.device(f"cuda:{cuda_device_id}")
    context = global_device_context_store.get_egl_context(device)
    if verbose:
        print(f"In thread {thread_id}, device {cuda_device_id}.")
    with context.active_and_locked():
        images = []
        for _ in range(3):
            images.append(_draw_square(**kwargs).float())
        if result is not None and thread_id is not None:
            egl_info = context.get_context_info()
            data = {"egl": egl_info, "images": images}
            result[thread_id] = data


@usesOpengl
class TestDeviceContextStore(TestCaseMixin, unittest.TestCase):
    def test_cuda_context(self):
        cuda_context_1 = global_device_context_store.get_cuda_context(
            device=torch.device("cuda:0")
        )
        cuda_context_2 = global_device_context_store.get_cuda_context(
            device=torch.device("cuda:0")
        )
        cuda_context_3 = global_device_context_store.get_cuda_context(
            device=torch.device("cuda:1")
        )
        cuda_context_4 = global_device_context_store.get_cuda_context(
            device=torch.device("cuda:1")
        )
        self.assertIs(cuda_context_1, cuda_context_2)
        self.assertIs(cuda_context_3, cuda_context_4)
        self.assertIsNot(cuda_context_1, cuda_context_3)

    def test_egl_context(self):
        egl_context_1 = global_device_context_store.get_egl_context(
            torch.device("cuda:0")
        )
        egl_context_2 = global_device_context_store.get_egl_context(
            torch.device("cuda:0")
        )
        egl_context_3 = global_device_context_store.get_egl_context(
            torch.device("cuda:1")
        )
        egl_context_4 = global_device_context_store.get_egl_context(
            torch.device("cuda:1")
        )
        self.assertIs(egl_context_1, egl_context_2)
        self.assertIs(egl_context_3, egl_context_4)
        self.assertIsNot(egl_context_1, egl_context_3)


@usesOpengl
class TestUtils(TestCaseMixin, unittest.TestCase):
    def test_load_extensions(self):
        # This should work
        _define_egl_extension("eglGetPlatformDisplayEXT", egl.EGLDisplay)

        # And this shouldn't (wrong extension)
        with self.assertRaisesRegex(RuntimeError, "Cannot find EGL extension"):
            _define_egl_extension("eglFakeExtensionEXT", egl.EGLBoolean)

    def test_get_cuda_device(self):
        # This should work
        device = _get_cuda_device(0)
        self.assertIsNotNone(device)

        with self.assertRaisesRegex(ValueError, "Device 10000 not available"):
            _get_cuda_device(10000)

    def test_egl_convert_to_int_array(self):
        egl_attributes = {egl.EGL_RED_SIZE: 8}
        attribute_array = _egl_convert_to_int_array(egl_attributes)
        self.assertEqual(attribute_array._type_, ctypes.c_int)
        self.assertEqual(attribute_array._length_, 3)
        self.assertEqual(attribute_array[0], egl.EGL_RED_SIZE)
        self.assertEqual(attribute_array[1], 8)
        self.assertEqual(attribute_array[2], egl.EGL_NONE)


@usesOpengl
class TestOpenGLSingleThreaded(TestCaseMixin, unittest.TestCase):
    def test_draw_square(self):
        context = EGLContext(width=MAX_EGL_WIDTH, height=MAX_EGL_HEIGHT)
        with context.active_and_locked():
            rendering_result = _draw_square().float()
            expected_result = torch.zeros(
                (MAX_EGL_WIDTH, MAX_EGL_HEIGHT, 3), dtype=torch.float
            )
            start_px = int(MAX_EGL_WIDTH / 4)
            end_px = int(MAX_EGL_WIDTH * 3 / 4)
            expected_result[start_px:end_px, start_px:end_px, 0] = 255.0
            expected_result[start_px:end_px, start_px:end_px, 2] = 255.0

        self.assertTrue(torch.all(expected_result == rendering_result))

    def test_render_two_squares(self):
        # Check that drawing twice doesn't overwrite the initial buffer.
        context = EGLContext(width=MAX_EGL_WIDTH, height=MAX_EGL_HEIGHT)
        with context.active_and_locked():
            red_square = _draw_square(r=1.0, g=0.0, b=0.0)
            blue_square = _draw_square(r=0.0, g=0.0, b=1.0)

        start_px = int(MAX_EGL_WIDTH / 4)
        end_px = int(MAX_EGL_WIDTH * 3 / 4)

        self.assertTrue(
            torch.all(
                red_square[start_px:end_px, start_px:end_px]
                == torch.tensor([255, 0, 0])
            )
        )
        self.assertTrue(
            torch.all(
                blue_square[start_px:end_px, start_px:end_px]
                == torch.tensor([0, 0, 255])
            )
        )


@usesOpengl
class TestOpenGLMultiThreaded(TestCaseMixin, unittest.TestCase):
    def test_multiple_renders_single_gpu_single_context(self):
        _draw_squares_with_context()

    def test_multiple_renders_single_gpu_context_store(self):
        _draw_squares_with_context_store()

    def test_render_two_threads_single_gpu(self):
        self._render_two_threads_single_gpu(_draw_squares_with_context)

    def test_render_two_threads_single_gpu_context_store(self):
        self._render_two_threads_single_gpu(_draw_squares_with_context_store)

    def test_render_two_threads_two_gpus(self):
        self._render_two_threads_two_gpus(_draw_squares_with_context)

    def test_render_two_threads_two_gpus_context_store(self):
        self._render_two_threads_two_gpus(_draw_squares_with_context_store)

    def _render_two_threads_single_gpu(self, draw_fn):
        result = [None] * 2
        thread1 = threading.Thread(
            target=draw_fn,
            kwargs={
                "cuda_device_id": 0,
                "result": result,
                "thread_id": 0,
                "r": 1.0,
                "g": 0.0,
                "b": 0.0,
            },
        )
        thread2 = threading.Thread(
            target=draw_fn,
            kwargs={
                "cuda_device_id": 0,
                "result": result,
                "thread_id": 1,
                "r": 0.0,
                "g": 1.0,
                "b": 0.0,
            },
        )

        thread1.start()
        thread2.start()
        thread1.join()
        thread2.join()

        start_px = int(MAX_EGL_WIDTH / 4)
        end_px = int(MAX_EGL_WIDTH * 3 / 4)
        red_squares = torch.stack(result[0]["images"], dim=0)[
            :, start_px:end_px, start_px:end_px
        ]
        green_squares = torch.stack(result[1]["images"], dim=0)[
            :, start_px:end_px, start_px:end_px
        ]
        self.assertTrue(torch.all(red_squares == torch.tensor([255.0, 0.0, 0.0])))
        self.assertTrue(torch.all(green_squares == torch.tensor([0.0, 255.0, 0.0])))

    def _render_two_threads_two_gpus(self, draw_fn):
        # Contrary to _render_two_threads_two_gpus, this renders in two separate threads
        # but on a different GPU each. This means using different EGL contexts and is a
        # much less risky endeavour.
        result = [None] * 2
        thread1 = threading.Thread(
            target=draw_fn,
            kwargs={
                "cuda_device_id": 0,
                "result": result,
                "thread_id": 0,
                "r": 1.0,
                "g": 0.0,
                "b": 0.0,
            },
        )
        thread2 = threading.Thread(
            target=draw_fn,
            kwargs={
                "cuda_device_id": 1,
                "result": result,
                "thread_id": 1,
                "r": 0.0,
                "g": 1.0,
                "b": 0.0,
            },
        )
        thread1.start()
        thread2.start()
        thread1.join()
        thread2.join()
        self.assertNotEqual(
            result[0]["egl"]["context"].address, result[1]["egl"]["context"].address
        )

        start_px = int(MAX_EGL_WIDTH / 4)
        end_px = int(MAX_EGL_WIDTH * 3 / 4)
        red_squares = torch.stack(result[0]["images"], dim=0)[
            :, start_px:end_px, start_px:end_px
        ]
        green_squares = torch.stack(result[1]["images"], dim=0)[
            :, start_px:end_px, start_px:end_px
        ]
        self.assertTrue(torch.all(red_squares == torch.tensor([255.0, 0.0, 0.0])))
        self.assertTrue(torch.all(green_squares == torch.tensor([0.0, 255.0, 0.0])))

    def test_render_multi_thread_multi_gpu(self):
        # Multiple threads using up multiple GPUs; more threads than GPUs.
        # This is certainly not encouraged in practice, but shouldn't fail. Note that
        # the context store will only allow one rendering at a time to occur on a
        # single GPU, even across threads.
        n_gpus = torch.cuda.device_count()
        n_threads = 10
        kwargs = {
            "r": 1.0,
            "g": 0.0,
            "b": 0.0,
            "verbose": True,
        }

        threads = []
        for thread_id in range(n_threads):
            kwargs.update(
                {"cuda_device_id": thread_id % n_gpus, "thread_id": thread_id}
            )
            threads.append(
                threading.Thread(
                    target=_draw_squares_with_context_store, kwargs=dict(kwargs)
                )
            )

        for thread in threads:
            thread.start()
        for thread in threads:
            thread.join()


@usesOpengl
class TestOpenGLUtils(TestCaseMixin, unittest.TestCase):
    @classmethod
    def tearDownClass(cls):
        global_device_context_store.set_context_data(torch.device("cuda:0"), None)

    def test_device_context_store(self):
        # Most of DCS's functionality is tested in the tests above, test the remainder.
        device = torch.device("cuda:0")
        global_device_context_store.set_context_data(device, 123)

        self.assertEqual(global_device_context_store.get_context_data(device), 123)

        self.assertEqual(
            global_device_context_store.get_context_data(torch.device("cuda:1")), None
        )

        # Check that contexts in store can be manually released (although that's a very
        # bad idea! Don't do it manually!)
        egl_ctx = global_device_context_store.get_egl_context(device)
        cuda_ctx = global_device_context_store.get_cuda_context(device)
        egl_ctx.release()
        cuda_ctx.detach()

        # Reset the contexts (just for testing! never do this manually!). Then, check
        # that first running DeviceContextStore.release() will cause subsequent releases
        # to fail (because we already released all the contexts).
        global_device_context_store._cuda_contexts = {}
        global_device_context_store._egl_contexts = {}

        egl_ctx = global_device_context_store.get_egl_context(device)
        cuda_ctx = global_device_context_store.get_cuda_context(device)
        global_device_context_store.release()
        with self.assertRaisesRegex(EGLError, "EGL_NOT_INITIALIZED"):
            egl_ctx.release()
        with self.assertRaisesRegex(pycuda._driver.LogicError, "cannot detach"):
            cuda_ctx.detach()

    def test_no_egl_error(self):
        # Remove EGL, import OpenGL with the wrong backend. This should make it
        # impossible to import OpenGL.EGL.
        del os.environ["PYOPENGL_PLATFORM"]
        modules = list(sys.modules)
        for m in modules:
            if "OpenGL" in m:
                del sys.modules[m]
        import OpenGL.GL  # noqa

        self.assertFalse(_can_import_egl_and_pycuda())

        # Import OpenGL back with the right backend. This should get things on track.
        modules = list(sys.modules)
        for m in modules:
            if "OpenGL" in m:
                del sys.modules[m]

        os.environ["PYOPENGL_PLATFORM"] = "egl"
        self.assertTrue(_can_import_egl_and_pycuda())

    def test_egl_release_error(self):
        # Creating two contexts on the same device will lead to trouble (that's one of
        # the reasons behind DeviceContextStore). You can release one of them,
        # but you cannot release the same EGL resources twice!
        ctx1 = EGLContext(width=100, height=100)
        ctx2 = EGLContext(width=100, height=100)

        ctx1.release()
        with self.assertRaisesRegex(EGLError, "EGL_NOT_INITIALIZED"):
            ctx2.release()