Расширенный пример для понимания CUDA, Numba, Cupy и т. д.

В основном все примеры Numba, CuPy и т. д., доступные в Интернете, представляют собой простые дополнения к массиву, демонстрирующие ускорение перехода от одиночного ядра/потока процессора к графическому процессору. А в документации по командам в основном отсутствуют хорошие примеры. Этот пост предназначен для предоставления более полного примера.

Исходный код предоставляется здесь. Это простая модель классических клеточных автоматов. Изначально он даже не использует numpy, а просто Python и модуль Pyglet для визуализации.

Моя цель — расширить этот код для конкретной проблемы (которая будет очень большой), но сначала я подумал, что лучше уже оптимизировать использование GPU.

Game_of_life.py таков:

import random as rnd
import pyglet
#import numpy as np
#from numba import vectorize, cuda, jit

class GameOfLife: 
 
    def __init__(self, window_width, window_height, cell_size, percent_fill):
        self.grid_width = int(window_width / cell_size) # cell_size 
        self.grid_height = int(window_height / cell_size) # 
        self.cell_size = cell_size
        self.percent_fill = percent_fill
        self.cells = []
        self.generate_cells()
  
    def generate_cells(self):
        for row in range(0, self.grid_height): 
            self.cells.append([])
            for col in range(0, self.grid_width):
                if rnd.random() < self.percent_fill:
                    self.cells[row].append(1)
                else:
                    self.cells[row].append(0)
                
    def run_rules(self): 
        temp = []
        for row in range(0, self.grid_height):
            temp.append([])
            for col in range(0, self.grid_width):
                cell_sum = sum([self.get_cell_value(row - 1, col),
                                self.get_cell_value(row - 1, col - 1),
                                self.get_cell_value(row,     col - 1),
                                self.get_cell_value(row + 1, col - 1),
                                self.get_cell_value(row + 1, col),
                                self.get_cell_value(row + 1, col + 1),
                                self.get_cell_value(row,     col + 1),
                                self.get_cell_value(row - 1, col + 1)])
                
                if self.cells[row][col] == 0 and cell_sum == 3:
                    temp[row].append(1)
                elif self.cells[row][col] == 1 and (cell_sum == 3 or cell_sum == 2):
                    temp[row].append(1)
                else:                 
                    temp[row].append(0)
        
        self.cells = temp

    def get_cell_value(self, row, col): 
        if row >= 0 and row < self.grid_height and col >= 0 and col < self.grid_width:
           return self.cells[row][col]
        return 0

    def draw(self): 
        for row in range(0, self.grid_height):
            for col in range(0, self.grid_width):
                if self.cells[row][col] == 1:
                    #(0, 0) (0, 20) (20, 0) (20, 20)
                    square_coords = (row * self.cell_size,                  col * self.cell_size,
                                     row * self.cell_size,                  col * self.cell_size + self.cell_size,
                                     row * self.cell_size + self.cell_size, col * self.cell_size,
                                     row * self.cell_size + self.cell_size, col * self.cell_size + self.cell_size)
                    pyglet.graphics.draw_indexed(4, pyglet.gl.GL_TRIANGLES,
                                         [0, 1, 2, 1, 2, 3],
                                         ('v2i', square_coords))

Во-первых, я мог бы использовать добавление numpy в конце generate_cells this self.cells = np.asarray(self.cells) и в конце run_rules this self.cells = np.asarray(temp), поскольку выполнение этого раньше не приведет к ускорению, как показано здесь. (На самом деле переход на numpy не привел к заметному ускорению)

Что касается графических процессоров, например, я добавлял @jit перед каждой функцией и стал очень медленным. Также пытался использовать @vectorize(['float32(float32, float32)'], target='cuda'), но возник вопрос: как использовать @vectorize в функциях, которые имеют только self в качестве входного аргумента?

Я также пытался заменить numpy на cupy, например self.cells = cupy.asarray(self.cells), но тоже стал очень медленным.

Следуя первоначальной идее расширенного примера использования графического процессора, каким будет правильный подход к проблеме? Где находится правильное место для размещения модификаций/векторизации/параллелизации/numba/cupy и т. д.? И самое главное, почему?

Дополнительная информация: помимо предоставленного кода, вот файл main.py:

import pyglet
from game_of_life import GameOfLife 
 
class Window(pyglet.window.Window):
 
    def __init__(self):
        super().__init__(800,800)
        self.gameOfLife = GameOfLife(self.get_size()[0],
                                     self.get_size()[1],
                                     15,  # the lesser this value, more computation intensive will be
                                     0.5) 

        pyglet.clock.schedule_interval(self.update, 1.0/24.0) # 24 frames per second
 
    def on_draw(self):
        self.clear()
        self.gameOfLife.draw()
        
    def update(self, dt):
        self.gameOfLife.run_rules()
 
if __name__ == '__main__':
    window = Window()
    pyglet.app.run()