Максимальная оптимизация игры «Жизнь» на Julia

7fb74fe3ad6a22820f5daccb37c8103d.jpg

Это очень хорошой case для оптимизации. Алгоритм крайне прост и его знают все. Но сколько можно сделать!

1. Julia, попытка первая и наивная

Julia хорошо изображает из себя питоно-подобный язык с утиной типизацией, будучи по реализации совершенно не питоно подобной. При том что синтаксис у нее отличается от Питона (и индексы начинаются с 1), почти каждую строчку надо переписывать, но все эти переписывания кажутся довольно тривиальными, что делает ее хорошим языком для проекта, когда Python уже жмет. Но эта похожесть может сыграть дурную службу.

Наша первая реализация:

global size = 10000
global board=zeros(Int8, size, size) 
global helper=zeros(Int8, size, size)

function step0()
  global size, board, helper
  for x in 1:size
    left = ((x == 1) ? size : x-1)
    lower = size
    right = ((x == size) ? 1 : x+1)
    for y in 1:size
      upper = ((y == size) ? 1 : y+1)
      # get the number of neighbours
      nb = board[left, upper] + board[x, upper] + board[right, upper] + 
                board[left, y] +                            board[right, y] + 
                board[left, lower] + board[x, lower]      + board[right, lower]
      cell = board[x, y]
      w = 0
      if (nb == 3) && (cell == 0) 
        w = 1 # new born
      elseif (nb < 2) && (cell == 1) 
        w = 1 # die lonely
      elseif (nb > 3) && (cell == 1) 
        w = 1 # die from owercrowding
      end
      helper[x, y] = w
      lower = y
    end
  end
  
  # now move modifications back
  Threads.@threads for x in 1:size
    for y in 1:size
	  if helper[x,y] > 0 
	     board[x,y] = 1 - board[x,y]
	  end
    end
  end
  
  return
end

Дает время 85.486411 seconds (1.02 G allocations: 18.188 GiB, 1.83% gc time). Что? Лишь немного быстрее Python?

Мы совершили страшную ошибку. Подсказкой является большое количество allocations. Дело в том, что в Julia переменные могут иметь любой тип. Поэтому, когда функция использует глобальные переменные, то она не может быть уверена в их типе, и вынуждена вставлять проверки run time. Правильным будет передавать все через параметры:

2. Первая правильная реализация

function step1(size, board, helper)
  for x in 1:size
    left = ((x == 1) ? size : x-1)
    lower = size
    right = ((x == size) ? 1 : x+1)
    for y in 1:size
      upper = ((y == size) ? 1 : y+1)
      # get the number of neighbours
      nb = board[left, upper] + board[x, upper] + board[right, upper] + 
                board[left, y] +                            board[right, y] + 
                board[left, lower] + board[x, lower]      + board[right, lower]
      cell = board[x, y]
      w = 0
      if (nb == 3) && (cell == 0) 
        w = 1 # new born
      elseif (nb < 2) && (cell == 1) 
        w = 1 # die lonely
      elseif (nb > 3) && (cell == 1) 
        w = 1 # die from owercrowding
      end
      helper[x, y] = w
      lower = y
    end
  end
  
  # now move modifications back
  for x in 1:size
    for y in 1:size
	  if helper[x,y] > 0 
	     board[x,y] = 1 - board[x,y]
	  end
    end
  end
  
  return
end

1.882592 seconds

Нет никаких аллокаций. Вот это уже другое дело. Мы достигли скорости 53M клеток в секунду (напоминаю, это AMD Ryzen 5 3600×6-Core Processor 3.79 GHz)

3. Multi-threading

К счастью, часто распараллелить программу на Julia элементарно — мы просто добавим Threads@threadsперед циклами.

function step2(size, board, helper)
  Threads.@threads for x in 1:size
    left = ((x == 1) ? size : x-1)
    lower = size
    right = ((x == size) ? 1 : x+1)
    for y in 1:size
      upper = ((y == size) ? 1 : y+1)
      # get the number of neighbours
      nb = board[left, upper] + board[x, upper] + board[right, upper] + 
                board[left, y] +                            board[right, y] + 
                board[left, lower] + board[x, lower]      + board[right, lower]
      cell = board[x, y]
      w = 0
      if (nb == 3) && (cell == 0) 
        w = 1 # new born
      elseif (nb < 2) && (cell == 1) 
        w = 1 # die lonely
      elseif (nb > 3) && (cell == 1) 
        w = 1 # die from owercrowding
      end
      helper[x, y] = w
      lower = y
    end
  end
  
  # now move modifications back
  Threads.@threads for x in 1:size
    for y in 1:size
	  if helper[x,y] > 0 
	     board[x,y] = 1 - board[x,y]
	  end
    end
  end
  
  return
end

Здесь мы пользуется тем, что циклы не пересекаются в памяти по записи. Тестировать будем на одном треде и с --threads 12:

1 thread: 1.981022 seconds (12 allocations: 1.375 KiB)

12 threads: 0.447547 seconds (124 allocations: 13.438 KiB)

Мы достигли скорости 223M клеток в секунду

4. inbounds, отмена проверки индексов массива

Так как мы абсолютно уверены в себе, давайте выключим проверки:

function step3(size, board, helper)
  Threads.@threads for x in 1:size
    left = ((x == 1) ? size : x-1)
    lower = size
    right = ((x == size) ? 1 : x+1)
    for y in 1:size
      upper = ((y == size) ? 1 : y+1)
      # get the number of neighbours
      @inbounds nb = board[left, upper] + board[x, upper] + board[right, upper] + 
                board[left, y] +                            board[right, y] + 
                board[left, lower] + board[x, lower]      + board[right, lower]
      @inbounds cell = board[x, y]
      w = 0
      if (nb == 3) && (cell == 0) 
        w = 1 # new born
      elseif (nb < 2) && (cell == 1) 
        w = 1 # die lonely
      elseif (nb > 3) && (cell == 1) 
        w = 1 # die from owercrowding
      end
      @inbounds helper[x, y] = w
      lower = y
    end
  end
  
  # now move modifications back
  Threads.@threads for x in 1:size
    for y in 1:size
	  @inbounds if helper[x,y] > 0 
	     @inbounds board[x,y] = 1 - board[x,y]
	  end
    end
  end
  
  return
end

1 thread: 1.584128 seconds (12 allocations: 1.375 KiB)

12 threads: 0.348541 seconds (127 allocations: 13.531 KiB)

Мы достигли скорости 286M клеток в секунду

5. Cache-friendly code

А что если мы поменяем везде индексы x, y на y, x?

function step4(size, board, helper)
  Threads.@threads for x in 1:size
    left = ((x == 1) ? size : x-1)
    lower = size
    right = ((x == size) ? 1 : x+1)
    for y in 1:size
      upper = ((y == size) ? 1 : y+1)
      # get the number of neighbours
      @inbounds nb = board[upper, left] + board[upper, x] + board[upper, right] + 
                board[y, left] +                            board[y, right] + 
                board[lower, left] + board[lower, x]      + board[lower, right]				
      @inbounds cell = board[y, x]
      w = 0
      if (nb == 3) && (cell == 0) 
        w = 1 # new born
      elseif (nb < 2) && (cell == 1) 
        w = 1 # die lonely
      elseif (nb > 3) && (cell == 1) 
        w = 1 # die from owercrowding
      end
      @inbounds helper[y, x] = w
      lower = y
    end
  end
  
  # now move modifications back
  Threads.@threads for x in 1:size
    for y in 1:size
	  @inbounds if helper[y,x] > 0 
	     @inbounds board[y,x] = 1 - board[y,x]
	  end
    end
  end
  
  return
end

1 threads: 0.842495 seconds (12 allocations: 1.375 KiB)

12 threads: 0.113791 seconds (125 allocations: 13.469 KiB)

Мы достигли скорости 879M клеток в секунду

Но почему так произошло? Дело в уровнях кеширования памяти. При обходе матрицы один порядок обхода монотонно увеличивает адрес доступа, тогда как другой читает байты «в разбивку». Как видно, на больших матрицах это очень существенно.

6. Убираем первый if

Отмена проверок границ массивов нам помогла не только потому, что часть кода оказалось выброшенной, но и потому, что мы избавились от if, а ветвления чреваты срывом конвейера. Если для проверки границ массивов хотя бы одна из веток почти никогда не выполняется, то вот этот код:

	  if helper[y,x] > 0 
	     board[y,x] = 1 - board[y,x]
	  end

ужасен. у этого IF нет 'правильной' стороны., но как вообще мне пришла в голову идея в helper записывать признак, что клетка поменялась? Просто я DBA и подсознательно оптимизировал объем UPDATE. А этого здесь делать не надо. В helper будем просто записывать новое состояние:

function step5(size, board, helper)
  Threads.@threads for x in 1:size
    left = ((x == 1) ? size : x-1)
    lower = size
    right = ((x == size) ? 1 : x+1)
    for y in 1:size
      upper = ((y == size) ? 1 : y+1)
      # get the number of neighbours
      @inbounds nb = board[upper, left] + board[upper, x] + board[upper, right] + 
                board[y, left] +                            board[y, right] + 
                board[lower, left] + board[lower, x]      + board[lower, right]				
      @inbounds cell = board[y, x]
      w = cell
      if (nb == 3) && (cell == 0) 
        w = 1 # new born
      elseif (nb < 2) && (cell == 1) 
        w = 0 # die lonely
      elseif (nb > 3) && (cell == 1) 
        w = 0 # die from owercrowding
      end
      @inbounds helper[y, x] = w
      lower = y
    end
  end
  
  # now move modifications back
  Threads.@threads for x in 1:size
    for y in 1:size
      @inbounds board[y,x] = helper[y,x]
    end
  end
  
  return
end

1 threads: 0.725642 seconds (12 allocations: 1.375 KiB)

12 threads: 0.089742 seconds (148 allocations: 16.094 KiB)

Мы достигли скорости 1.114B (миллиардов) клеток в секунду. Рубеж миллиарда пройден!

7. Убираем главный IF

Раз это так помогло, то давайте избавимся от главного IF, заменив его lookup — таблицей. Заодно программа станем много короче:

function step6(size, board, helper)
  #                                no cell  born           | live   2  3   4
  lookup = zeros(Int8, 18).= [0, 0, 0, 1, 0, 0, 0, 0, 0,      0, 0, 1, 1, 0, 0, 0, 0, 0]
  Threads.@threads for x in 1:size
    left = ((x == 1) ? size : x-1)
    lower = size
    right = ((x == size) ? 1 : x+1)
    for y in 1:size
      upper = ((y == size) ? 1 : y+1)
      # get the number of neighbours
      @inbounds nb = board[upper, left] + board[upper, x] + board[upper, right] + 
                board[y, left] +                            board[y, right] + 
                board[lower, left] + board[lower, x]      + board[lower, right]				
      @inbounds cell = board[y, x]
	  @inbounds helper[y, x] = lookup[ cell*9 + nb+1]
      lower = y
    end
  end
  
  # now move modifications back
  Threads.@threads for x in 1:size
    for y in 1:size
      @inbounds board[y,x] = helper[y,x]
    end
  end
  
  return
end

1 threads: 0.256231 seconds (14 allocations: 1.656 KiB)

12 threads: 0.053782 seconds (151 allocations: 16.406 KiB)

Мы достигли скорости 1.859B (миллиардов) клеток в секунду

8. Избавляемся от последних IF

У нас остались IF в виде? выражений для проверки границ тора. В принципе, эти условия можно переписать как арифметические операции с модулями, но увы, это лишь замедляет программу. Поэтому надо либо считать, что границы засеены нулями и считать все, кроме границы, либо написать специальный код для расчета границ.

Для краткости кода я остановился на первом решении:

function step7(size, board, helper)
  #                                no cell  born           | live   2  3   4
  lookup = zeros(Int8, 18).= [0, 0, 0, 1, 0, 0, 0, 0, 0,      0, 0, 1, 1, 0, 0, 0, 0, 0]
  sizem1 = size - 1
  Threads.@threads for x in 2:sizem1
    left = x-1
    right = x+1
    for y in 2:sizem1
      upper = y+1
      lower = y-1
      # get the number of neighbours
      @inbounds nb = board[upper, left] + board[upper, x] + board[upper, right] + 
                board[y, left] +                            board[y, right] + 
                board[lower, left] + board[lower, x]      + board[lower, right]				
      @inbounds cell = board[y, x]
	  @inbounds helper[y, x] = lookup[ cell*9 + nb+1]
    end
  end
  
  # now move modifications back
  Threads.@threads for x in 1:size
    for y in 1:size
      @inbounds board[y,x] = helper[y,x]
    end
  end
  
  return
end

1 threads: 0.164568 seconds (14 allocations: 1.656 KiB)

12 threads: 0.036595 seconds (151 allocations: 16.406 KiB)

Мы достигли скорости 2.732B (миллиардов) клеток в секунду

9. Векторные операции

Импортируем пакет LoopVectorization, в котором макро@turboпереписывает код цикла так, чтобы использовались векторные операции. При этом важно:

  • В цикле не должно быть никаких if (в том числе?)

  • Выполнение цикла для разных элементов не может зависеть друг от друга (например, там не может быть нарастающего итога)

  • Возможны любые сюрпризы

  • Так как любые IF запрещены, то@inboundsподразумевается всегда

function step8(size, board, helper)
  lookup = zeros(Int8, 18).= [0, 0, 0, 1, 0, 0, 0, 0, 0,      0, 0, 1, 1, 0, 0, 0, 0, 0]
  sizem1 = size - 1
  Threads.@threads for x in 2:sizem1
    left = x-1
    right = x+1
    @turbo for y in 2:sizem1
      upper = y+1
      lower = y-1
      nb = board[upper, left] + board[upper, x] + board[upper, right] + 
                board[y, left] +                            board[y, right] + 
                board[lower, left] + board[lower, x]      + board[lower, right]				
      cell = board[y, x]
	  helper[y, x] = lookup[ cell*9 + nb+1]
    end
  end
  
  Threads.@threads for x in 1:size
    @turbo for y in 1:size
      board[y,x] = helper[y,x]
    end
  end
end

1 threads: 0.069448 seconds (14 allocations: 1.656 KiB)

12 threads: 0.021376 seconds (154 allocations: 16.500 KiB)

Но увы! Результат оказывается неправильным. Все таки первое выражение оказывается слишком сложным для векторизации. Второе@turboработает корректно, но не дает прироста производительности. Однако сама идея 'переписывающих макросов' замечательная!

10. Уменьшение чтений

Как видно, каждая ячейка читается 9 раз. Можно уменьшить это количество до трех, организовав своебразный 'конвейер' из ячеек, смещая их на кждом шаге и читая только значения с одной стороны:

function step9(size, board, helper)
  lookup = zeros(Int8, 18).= [0, 0, 0, 1, 0, 0, 0, 0, 0,   0, 0, 1, 1, 0, 0, 0, 0, 0]
  sizem1 = size - 1
  Threads.@threads for x in 2:sizem1
    left = x-1
    right = x+1
	cell_left_lower = cell_lower = cell_right_upper = 0
	@inbounds cell_left = board[2, left]
	@inbounds cell_right = board[2, right]
	cell = board[2, x]
    @inbounds for y in 2:sizem1
      upper = y+1
      lower = y-1
	  cell_left_upper = board[upper, left]
	  cell_right_upper = board[upper, right]
	  cell_upper = board[upper, x]
      nb = cell_left_upper + cell_upper + cell_right_upper + 
             cell_left +           cell_right + 
             cell_left_lower + cell_lower + cell_right_upper				
	  helper[y, x] = lookup[ cell*9 + nb+1]
	  # shift
	  cell_left_lower = cell_left
	  cell_lower = cell
	  cell_right_lower = cell_right
	  cell_left = cell_left_upper
	  cell = cell_upper
	  cell_right = cell_right_upper
    end
  end
  
  Threads.@threads for x in 1:size
    for y in 1:size
      @inbounds board[y,x] = helper[y,x]
    end
  end
end

Как ни странно. стало хуже:

1 threads: 0.188678 seconds (14 allocations: 1.656 KiB)

12 threads: 0.040889 seconds (152 allocations: 16.438 KiB)

А у вас есть идеи, почему стало хуже?

11. Храним 8 клеток в байте

Если мы хотим обработать сразу 8 бит, то итоговое состояние зависит от : данных 8 клеток, 8 клеток над и под ними (итого 24), и дополнительных границ слева и справа по три клетки, что дает 30 бит и lookup таблицу в 1Gb — не так много по современным меркам, но потребуется битовая магия:

global size = 10000 # x
global size8 = trunc(Int, size/8) # y
global board=zeros(UInt8, size8, size) 
global helper=zeros(UInt8, size8, size)
global lookup=zeros(UInt8, 256*256*256*64)

unction step10(size8, size, board, helper, lookup)
  sizem1 = size - 1
  size8m1 = size8 - 1
  Threads.@threads for x in 2:sizem1
    left = x-1
    right = x+1
	cell_left_lower = cell_lower = cell_right_upper = cell_right_lower = 0
	@inbounds cell_left = board[2, left]
	@inbounds cell_right = board[2, right]
	cell = board[2, x]
    @inbounds for y in 2:size8m1
      upper = y+1
      lower = y-1
	  cell_left_upper = board[upper, left]
	  cell_right_upper = board[upper, right]
	  cell_upper = board[upper, x]

	  helper[y, x] = lookup[cell | (cell_upper<<8) | (cell_lower<<16) +
	    ((cell_left_upper&1)<<24) | 
        ((cell_left&1)<<25) |
        ((cell_left_lower&1)<<26) |      # lower bit, &1
	    ((cell_right_upper&0x80)<<20) | ((cell_right&0x80)<<21) 
        | ((cell_right_lower&0x80)<<22) # high bit, &127
		]
	  # shift
	  cell_left_lower = cell_left
	  cell_lower = cell
	  cell_right_lower = cell_right
	  cell_left = cell_left_upper
	  cell = cell_upper
	  cell_right = cell_right_upper
    end
  end
  
  Threads.@threads for x in 1:size
    for y in 1:size8
      @inbounds board[y,x] = helper[y,x]
    end
  end
end

1 threads: 0.015258 seconds (12 allocations: 1.375 KiB)

12 threads: 0.003306 seconds (149 allocations: 16.125 KiB)

30.248B клеток в секунду. Вау.

12. 32 бита сразу.

Мы можем попробовать обрабатывать сразу 32 бита, уйдя в совсем глубокую битовую магию:

      helper[y,x] = 	  
  	    lookup[(cell&0xFF) | ((cell_upper&0xFF)<<8) | ((cell_lower&0xFF)<<16) +
	      ((cell_left_upper&1)<<24) | ((cell_left&1)<<25) | ((cell_left_lower&1)<<26) |  
	      ((cell_upper&0x100)<<19) | ((cell&0x100)<<20) | ((cell_lower&0x100)<<21) 
	   	  ] |
  	    (lookup[(cell&0xFF00>>8) | ((cell_upper&0xFF00)) | ((cell_lower&0xFF00)<<8) +
	      ((cell_upper&0x80)<<17) | ((cell_left&0x80)<<18) | ((cell_lower&0x80)<<19) |  
	      ((cell_upper&0x10000)<<12) | ((cell&0x10000)<<13) | ((cell_lower&0x10000)<<14) 
	   	  ] >> 8) |
  	    (lookup[(cell&0xFF0000>>16) | ((cell_upper&0xFF0000)>>8) | ((cell_lower&0xFF0000)) +
	      ((cell_upper&0x8000)<<9) | ((cell_left&0x8000)<<10) | ((cell_lower&0x8000)<<11) |  
	      ((cell_upper&0x1000000)<<4) | ((cell&0x1000000)<<5) | ((cell_lower&0x1000000)<<6) 
	   	  ] >> 16) |
  	    (lookup[(cell&0xFF000000>>24) | ((cell_upper&0xFF000000)>>16) | ((cell_lower&0xFF000000)>>8) +
	      ((cell_upper&0x800000)<<1) | ((cell_left&0x800000)<<2) | ((cell_lower&0x800000)<<3) |  
	      ((cell_right_upper&0x80000000)>>4) | ((cell_right&80000000)>>3) | ((cell_right_lower&80000000)>>2) 
	   	  ] >> 24) 

Но это примерно в 10 раз медленнее…Слищком много битовых операций для упаковки/распаковки…

© Habrahabr.ru