Перенос молекулярной динамики на CUDA. Часть I: Основы

Цель данной статьи – поднять вопросы распараллеливания кода программы для численного моделирования методом молекулярной динамики (МД) с помощью технологии CUDA. Зачем это вообще нужно, ведь уже существуют программные пакеты по МД, работающие в том числе и на CUDA? Дело в том, что я развиваю свою собственную концепцию «непостоянного поля сил» (non-constant force field), которая не реализована в существующих МД-программах.

Переделывать чужой код под эти нужды – довольно неблагодарное занятие, поэтому я взялся перенести уже написанный свой последовательный код и заодно поделится некоторыми размышлениями. Кроме того, это ответ на часто мелькающий здесь комментарий к статьям по CUDA, вроде этого .

Итак, что же такое молекулярная динамика? На Хабре уже есть несколько постов на эту тему, например здесь или вот здесь. Кратко, МД – это метод, позволяющий моделировать движение множества частиц (в том числе атомов, ионов, молекул) и рассчитывать коллективные свойства системы, зависящие от этого движения. Как это работает? Допустим для множества из N частиц заданы некоторые начальные координаты, скорости, массы и (главное!) законы взаимодействия между ними. Изменяем координаты согласно скоростям. На основе законов взаимодействия вычисляем силы, действующие между частицами. Раз знаем силу и массу – знаем ускорение. Поправляем скорость с учетом ускорения. И снова переходим к изменению координат. И так повторяем тысячи раз, пока не надоест не наберем достаточную статистику.



В реальных приложениях последовательность вычислений «координаты – скорости – силы» несколько сложнее, основаны на разложении в ряд функции координаты частицы. Наиболее распространенным является скоростной алгоритм Верле (Velocity Verlet). Псевдокод МД программы на основе интегратора Верле выглядит следующим образом:

init_md;
int iStep = 0;
while (iStep < nStep)
{
    reset_quantities;
    clear_cell_list;
    verlet_1stage;
    forces_calculation;
    verlet_2stage;
    temp_scale;
    termostat;
    some_output;
    iStep++;
}
final_md; 

Основной цикл программы выполняется заданное пользователем число раз (nStep). Итерация этого цикла называется шагом по времени или просто шагом (timestep). Функции init_md, reset_quantites, some_output и final_md мы рассматривать не будем, они считывают входную информацию, обнуляют величины, выводят промежуточные данные, выделяют и высвобождают память и т.д. С точки зрения вывода информации, молекулярная динамика идеальна для проведения расчетов на видеокарте: большую часть вычислений вообще нет необходимости выгружать.

Функции verlet_1stage, verlet_2stage, temp_scale и termostat относятся к интегратору уравнений движения и учету температуры. Функция forces_calculation вычисляет силы, действующие между частицами, а функция clear_cell_list относится к технике ускорения вычислений (об этом подробнее далее).

Рассмотрим все эти группы функций.

1 Интегрирование уравнений движения

Суть первой стадии алгоритма Верле можно выразить следующими формулами:

Второй:

где X – координата частицы, V – её скорость, F – сила действующая на неё, m – её масса и dt – шаг по времени. Все это применяется ко всем частицам и, естественно, происходит в 3х измерениях, итак:

__global__ void verlet_1stage(int atPerBlock, int atPerThread, cudaMD *md)
{
  int i, t
  float charge;
  float engElecField = 0.0;		// энергия во внешнем электрическом поле
  __shared__ float shEngElField;	// shared версия этой переменной

  // сбросим энергию в shared переменных
    if (threadIdx.x == 0)
        shEngElField = 0;
    __syncthreads();

  // определим диапазон частиц, с которыми будет работать поток
 int id0 = blockIdx.x * atPerBlock + threadIdx.x * atPerThread;
 int N = min(id0 + atPerThread, md->nAt);

  for (i = id0; i < N; i++)
  {

     t = md->types[i];
     charge = md->specs[t].charge;

     //  v = v + f/m * 0.5 dt
     md->vls[i].x += md->rMasshdT[i] * md->frs[i].x;
     md->vls[i].y += md->rMasshdT[i] * md->frs[i].y;
     md->vls[i].z += md->rMasshdT[i] * md->frs[i].z;

     // x = x + v * dt
     md->xyz[i].x += md->vls[i].x * md->tSt;
     md->xyz[i].y += md->vls[i].y * md->tSt;
     md->xyz[i].z += md->vls[i].z * md->tSt;

     put_periodic(&md->xyz[i], md->vls[i], md->masses[i], md->types[i], md);

     keep_in_cell(i, md->xyz[i], md);

     // внешнее электрическое поле
     engElecField += charge * md->xyz[i].x * md->elecField.x; // Eng = q * x * dU/dx
  }

  // добавляем энергию поля в shared перменную…
  atomicAdd(&shEngElField, engElecField);
   __syncthreads();

  //.. а затем в глобальную
  if (!threadIdx.x)   // 0th thread
  {
      atomicAdd(&md->engElecField, shEngElField);
  }
}

Входные параметры atPerBlock и atPerTherad равны количеству частиц (атомов) обрабатываемых одним блоком и потоком, соответсвенно. Поскольку число атомов в таких расчетах обычно постоянно, эти переменные можно вычислить раз и навсегда где-нибудь в функции init_md. Параметр md типа *cudaMD – ссылка на структуру, где содержаться ссылки на все необходимые для расчета переменные, такие как:

• nAt – число частиц;
• xyz – массив координат типа float3;
• vls – массив скоростей типа float3 (3 компонента вектора);
• frs – массив сил (также 3 компонета вектора);
• массив rMasshdT хранит сразу выражение 0.5*dt*m для каждой частицы;
• types – массив типов частиц, ведь атомы могут быть разные;
• массив spec хранит свойства частиц, массу, заряд и т.д. Чтобы получить свойства i-ой частицы нужно сделать обращение вида spec[types[i]].

Я объединил это всё в структуру, чтобы не передавать в функцию кучу параметров. Возможно, это не лучшее решение? Может быть стоило объявить их глобальными переменными? Не в том смысле, что они находятся в global memory, а в том, что они видны во всех функциях программы. Функцию keep_in_cell обсудим в разделе ускорения вычислений, а функция put_periodic помещает частицу в моделируемый объём (бокс, box), если она вылетит за его пределы (как в пакмане, как метко выразился Dare в своем посте про МД). Ну и раз уж описанная функция выполняет обновление координат частиц, то имеет смысл посчитать все энергии, которые зависят только от координат, например энергию частиц во внешнем поле, если таковое было задано. Здесь и далее, по справделивому замечанию avdx, я убрал из кода функциq убрал обнуление таких переменных как md->engElecField и md->engKin, для этих целей теперь выделена функция reset_quantites. Вопрос по функции verlet_1stage следующий: некоторые переменные (координаты, скорости и т.д.) используются по нескольку раз и целесообразно ли их сначала выгрузить в какую-то более быструю память? Или компилятор сам умеет оптимизировать подобные вещи?

Функция put_periodic делает так, чтобы частица, вылетевшая за границы моделируемой области (бокса) появилась с другой её стороны (периодические граничные условия):

__device__ void put_periodic(float3& xyz, float3 vel, float mass, int type, cudaMD *md)
{
   if (xyz.x < 0)
   {
      xyz.x += md->leng.x;

      atomicAdd(&md->specAcBoxNeg[type].x, 1); 
      atomicAdd(&md->negMom.x, mass * (-vel.x)); 
   }
   else
     if (xyz.x >= md->leng.x)
     {
        xyz.x -= md->leng.x;

        atomicAdd(&md->specAcBoxPos[type].x, 1);
        atomicAdd(&md->posMom.x, mass * vel.x);
     }
}

Функция работает для бокса с геометрией прямоугольного параллелепипеда. Код приведен для измерения x, остальные измерения полностью аналогичны. Поля leng и revLeng – размер бокса в соответствующем измерении и обратная от него величина. Вообще для возвращения частицы в бокс можно обойтись без условных операторов:

xyz.x -= floor(xyz.x * md->revLeng.x) * md->leng.x;

Я использовал условные операторы в последовательной версии, поскольку мне надо было считать количество частиц, прошедших через границу бокса и суммарный импульс этих частиц (из него можно выразить давление). В приведенном коде также включены эти счетчики, возможно для CUDA надо будет подумать, как обойтись без if.

Вторая часть скоростного интегратора Верле:

__global__ void verlet_2stage(int atPerBlock, int atPerThread, cudaMD *md)
{
  int i;
  float kinE = 0.0;		// кинетическая энергия
  __shared__ float shKinE;	// shared версия

 int id0 = blockIdx.x * atPerBlock + threadIdx.x * atPerThread;
 int N = min(id0 + atPerThread, md->nAt);

  // обнулим кинетическую энергию в shared переменных
 if (threadIdx.x == 0)
        shKinE = 0.0;
  __syncthreads();


  // v = v + f/m * 0.5 dt
  for (i = id0; i < N; i++)
  {
    md->vls[i].x += md->rMasshdT[i] * md->frs[i].x;
    md->vls[i].y += md->rMasshdT[i] * md->frs[i].y;
    md->vls[i].z += md->rMasshdT[i] * md->frs[i].z;

    // сохраняем mv2
    kinE += (md->vls[i].x * md->vls[i].x + md->vls[i].y  * md->vls[i].y + md->vls[i].z * md->vls[i].z) * md->masses[i];
  }

  // суммируем кинетическую энергию сначала в shared memory...
  atomicAdd(&shKinE, 0.5 * kinE);   // kinE = mv2/2
  __syncthreads();
  //... затем в global memory
  if (!threadIdx.x)
     atomicAdd(&md->engKin, shKinE);
}

Код по структуре полностью аналогичен первой части интегратора. Обновляем скорости для каждой частицы, суммируем кинетическую энергию.

«Температурные» функции temp_scale и termostat. Молекулярные системы отличаются от механических тем, что в них есть понятие температуры. МД унаследовала определение температуры из молекулярно-кинетической теории, где температура – величина пропорциональная среднекинетической энергии частиц (есть основания полагать, что это заблуждение). Для поддержания нужной температуры в таких терминах можно просто масштабировать скорости частиц, это делает функция temp_scale:

__global__ void temp_scale(int atPerBlock, int atPerThread, cudaMD *md)
{
   float k;
   int i;
   int id0 = blockIdx.x * atPerBlock + threadIdx.x * atPerThread;
   int N = min(id0 + atPerThread, md->nAt);

   // вычислить предварительно
   //md->k = sqrt(md->teKin / md->engKin);
   for (i = 0; i < N; i++)
   {
      md->vls[i].x *= md->k;
      md->vls[i].y *= md->k;
      md->vls[i].z *= md->k;
   }

   if (!blockIdx.x)
    if (!threadIdx.x)
     {
        md->engKin = md->teKin;
     }
} 

Поле teKin содержит целевую кинетическую энергию, она пропорциональна температуре, установленной пользователем. Функция просто домножает все компоненты векторов скоростей на величину k, а затем один поток устанавливает значение кинетической энергии равной заданной. Переменную k необходимо определить перед вызовом данной функции, например в reset_quantities. Такое примитивное масштабирование температуры применяется обычно на предварительных этапах моделирования (equilibration period) с частотой раз в несколько шагов. Более хитрые алгоритмы поддержания температуры (термостаты) имитируют столкновения частиц с некоторым резервуаром заданной температуры, наиболее популярный – термостат Нозе-Гувера:

__global__ void tstat_nose(int atPerBlock, int atPerThread, cudaMD *md)
{
   int i;

//   этот код должен быть выполнен в отдельной функции
//   if (!blockIdx.x)
//     if (!threadIdx.x)   
//        md->chit += md->tSt * (md->engKin - md->teKin) * md->rQmass;
//   __synctreads();

   //вычислить предварительно:
   //float md->scale = 1.0 - md->tSt * md->chit;

   int id0 = blockIdx.x * atPerBlock + threadIdx.x * atPerThread;
   int N = min(id0 + atPerThread, md->nAt);

   for (i = id0; i < N; i++)
   {
       md->vls[i].x *= md->scale;
       md->vls[i].y *= md->scale;
       md->vls[i].z *= md->scale;
   }

   // вычислить после выполнения функции:
   //if (!blockIdx.x)
   //  if (!threadIdx.x)   
   //  {
   //     md->engKin = md->engKin * scale * scale;
   //     md->consInt += md->tSt * md->chit * md->qMassTau2;
   //     md->chit += md->tSt * (md->engKin - md->teKin) * md->rQmass; 
   //  }
}

Отличие от предыдущей функции лишь в более сложном определении поправочного коэффициента.

Подведем итог касательно этой группы функций. Все они прекрасно параллелизуются и не только потому, что выполняют абсолютно одинаковые действия ко всем частицам, но и потому, что не требуют перекрёстных ссылок. Поток, обрабатывая массивы с индексом i, не обращается к индексу j, нет конфликта доступа к одним и тем же данным.

2 Вычисление сил и список ячеек

В молекулярных системах основной источник сил, действующий на частицу (кроме уже упомянутых внешних полей) – это другие частицы. Вычисление этих сил и есть наиболее сложная, трудоемкая и неоднозначная часть МД. Одно из простейших приближений заключается в том, что все взаимодействия в системе представляют суммой взаимодействий между КАЖДОЙ парой частиц (парное приближение). Для этих целей вводят парный потенциал – функциональную зависимость потенциальной энергии взаимодействия двух частиц от расстояния между ними. Наиболее известные парные потенциалы: потенциал Леннард-Джонса (или «6-12»), Букингема и Борн-Майера-Хаггинса. Для примера приведу аналитические выражения первых двух:

и

где U – потенциальная энергия пары частиц, r – расстояние между частицами, а все остальные буквы – параметры, которые зависят только от типа взаимодействующих частиц и заданы на этапе инициализации. Однако даже больше, чем энергия нас интересуют силы взаимодействия между этими частицами. Силы находятся тривиально, через производную парного потенциала, для которой также несложно получить аналитическое выражение и зашить в программу. И, поскольку сила действия равна силе противодействия, мы сразу получаем силы, действующие на обе частицы со стороны друг друга, псевдокод:

U = 0;
for (i = 0; i < N; i++)
	for (j = i + 1; j < N-1; j++)
	{
		dx = x[i] – x[j];
		dy = y[i] – y[j];
		dz = z[i] – z[j];
		r = sqrt(dx*dx + dy*dy + dz*dz);
		U += pair_pot(r, types[i], types[j]);
		F = force(r, types[i], types[j])
		fx[i] = F * dx;
		fx[j] = -F * dx;
		fy[i] = F * dy;
		fy[j] = -F * dy;
		fz[i] = F * dz;
		fz[j] = -F * dz;

	}

Код простейший: находим расстояние между частицами, подставляем его в функцию для парного потенциала и в функцию для силы (обе зависят от типов взаимодействующих частиц) и вычисляем проекции сил. Цикл выполняется N(N-1)/2 раз, где N – число частиц. Чтобы хоть как-то сократить число вычислений можно воспользоваться тем, что величина данных потенциалов быстро убывает с расстоянием. На расстояниях больше 6-8 ангстрем, потенциал можно не вычислять, он будет почти равен нулю. Добавив условие проверки расстояния (а лучше его квадрата) мы сократим число вычислений, но не сократим число итераций цикла. Для больших систем большинство проверок будут давать false. Следующие техники позволяют уменьшить число проверок за счет отброса части пар с ЗАВЕДОМО большим межатомным расстоянием.

Существуют два основных метода сокращения числа переборов пар частиц: cписок соседей (он же список Верле, Verlet list) и список ячеек (cell list). В первом мы для каждой частицы запоминаем её соседей, которые находятся в радиусе действия потенциала. Список соседей необходимо обновлять с некоторой, заранее неизвестной, периодичностью, в чем и заключается недостаток метода. Поэтому остановимся на списке ячеек, суть которого заключается в разбиении моделируемого объёма (бокса) на (суб)ячейки (cell). Первоначально, размер ячейки выбирали больше либо равным радиусу обрезания. Тогда достаточно перебрать пары частиц внутри каждой ячейки и между соседними ячейками. Размер ячейки может быть и другим, главное определить какие пары ячеек надо проверить. Фактически это попытка описать сферу кубиками. Традиционно, для реализации списка ячеек в последовательном коде вводятся два массива, head и list. Размер первого равен числу ячеек, второго – числу частиц. В первом хранится индекс одной из частиц в заданной ячейке, индексы второго соответствуют индексам частиц, а содержимое – индексу следующей частицы в ячейке (-1, если частиц в ячейке больше нет), иллюстрация дана ниже.



Псевдокод обхода такого списка внутри ячейки с индексом cell_index:

i = head[cell_index];
while (i >= 0)
{
      j = list[i];
      while (j >= 0)
      {
           обрабатываем пару частиц i-j;
           j = list[j];
      }
i = list[i]; 
}	

Теперь попробуем реализовать список ячеек на CUDA. В нулевом приближении просто раздадим каждому потоку диапазон ячеек, в остальном оставив последовательный код перебора пар частиц внутри ячейки и пар с участием частиц соседних ячеек. Для каждой ячейки должен быть определен массив соседних ячеек lstHNeig[iC][jC] и число этих соседей, массив nHeadNeig[]:

__global__ void cell_list0(int cellPerBlock, int cellPerThread, cudaMD *md)
{
   int iC, jC;	// индексы ячеек
   int i, j;	// индексы частиц

   int id0 = blockIdx.x * cellPerBlock + threadId.x * cellPerThread;
   int N = min(id0 + cellPerThread, md->nCell);

   // цикл по ячейкам
   for (iC = id0; iC < N; iC++)
   {
      i = md->chead[iC];

      while (i >= 0)
      {
         // цикл внутри одной ячейки
         j = md->clist[i]; 	
         while (j >= 0)
         {
            обрабатываем пару частиц i-j;
            j = md->clist[j];  
         }

         // цикл по соседям ячейки
         for (jC = 0; jC < md->nHeadNeig[iC]; jC++)
         {
            j = md->chead[md->lstHNeig[iC][jC]];
            while (j >= 0)
            {
               обрабатываем пару частиц i-j;    
               j = md->clist[j];
            }
         }
         i = md->clist[i];
      } 
   }
}

На мой взгляд, такая реализация не раскрывает полностью потенциал видеокарты, поскольку один поток обрабатывает одну ячейку, которых как правило существенно меньше числа частиц. Во-вторых, внутри функции обработки пары частиц возможна работа только через atomic операции. Попробуем сильнее распараллелить задачу, пусть каждый поток обрабатывает пару соседних ячеек (считаем, что список таких пар уже создан):

__global__ void cell_list1(int pairPerBlock, int pairPerThread, cudaMD *md)
{
   int iP, iC, jC;	// индексы пары ячеек и ячеек
   int i, j;		// индексы частиц

   int id0 = blockIdx.x * pairPerBlock + threadId.x * pairPerThread;
   int N = min(id0 + cellPerThread, md->nCell);

   for (iP = id0; iC < N; iP++)
   {
      iC = md->cellPairs[iP].x;
      jC = md->cellPairs[iP].y;
      i = md->chead[iC];

      while (i >= 0)
      {
         j = md->chead[jC];
         while (j >= 0)
         {
             обрабатываем пару частиц i-j;
             j = md->clist[j];
         }
         i = md->clist[i];
      } 
   }
} 

Массив cellPairs[] содержит количество пар соседних ячеек (.x – первая ячейка пары, .y – вторая). Код проще и изящнее, но здесь пропадает обработка пар частиц внутри одной ячейки, для чего должна быть написана отдельная функция. Есть ещё одна проблема, организация списка ячеек и его обход неэффективны для параллельных вычислений, поскольку чтобы получить индекс следующей частицы в ячейке нужно пробежать по всем предыдущим. Попробуем создать более эффективный код и структуру данных, чтобы пару ячеек могли обрабатывать сразу несколько потоков. Пусть теперь список ячеек будет хранится в двумерном массиве cell[][], где первый индекс определяет номер ячейки, а второй – порядковый номер атома внутри ячейки. Также я храню число частиц в ячейке в нулевом элементе массива: cell[i][0], значение cell[i][j+1] равно общему индексу частицы в global memory. Правда такая реализация требуется знать максимальное количество частиц в ячейке и будет занимать больше памяти. Также разобьём нашу задачу на 2 подзадачи. Первая функция будет вычислять взаимодействие в непересекающихся парах ячеек (например, с индексами 0-1, 2-3, 4-5 и т.д.) и внутри ячеек:

__global__ void cell_list2a(cudaMD *md)
{
   int iP, iC, iCA, iCB, step, step1; 
   int i, j, id0, id1, N, N1, deltId;
   int id;
   int nA, nB, nOwn;  
   int ex = 0;		//флаг выхода из цикла

   // накопители Ван-дер-Ваальсовой и Кулоновской энергий внутри потока и в блоке
   __shared__ float engVdW, engCoul;
   float eVdW = 0.0;
   float eCoul = 0.0;

// динамически выделенная shared память
  extern __shared__ int shMem[];      
  int* sAids = shMem;		// номера частиц в ячейке 1 (тут она обозначена А)
  int* sAtypes = (int*)&sAids[md->maxAtPerCell]; // типы частиц в ячейке А
  float3* sAxyz = (float3*)&sAtypes[md->maxAtPerCell]; // координаты частиц в ячейке А
  float3* sAfrs = (float3*)&sAxyz[md->maxAtPerCell]; // силы на частице в ячейке А
   // аналогичные массивы для ячейки 2 (B)
  int* sBids = (int*)&sAfrs[md->maxAtPerCell];
  int* sBtypes = (int*)&sBids[md->maxAtPerCell];
  float3* sBxyz = (float3*)&sBtypes[md->maxAtPerCell];
  float3* sBfrs = (float3*)&sBxyz[md->maxAtPerCell];

   // указатели на массивы, которые обрабатывает данный поток
   float3 *sXyz;
   float3 *sFrs;
   int *sTypes;
   int *sIds;

   // берем индексы пар ячеек в данной паре
   iCA = md->cellPairs[blockIdx.x].x;  // индекс ячейки А
   iCB = md->cellPairs[blockIdx.x].y;  // индекс ячейки B

   // обнуляем Кулоновскую и Ван-дер-Ваальсову энергии
  if (threadIdx.x == 0)
    {
        engVdW = 0.0;
        engCoul = 0.0;
    }

   // сперва потоки копируют данные в shared_memory
   if (threadIdx.x < blockDim.x / 2) //! половина потоков будет копировать в ячейку A, половина - в ячейку B. Надо, чтобы размер блока был четным
     {
        iC = iCA;
        sXyz = sAxyz;
        sFrs = sAfrs;
        sTypes = sAtypes;
        sIds = sAids;
        deltId = 0;
     }
   else
     {
        iC = iCB;
        sXyz = sBxyz;
        sFrs = sBfrs;
        sTypes = sBtypes;
        sIds = sBids;
        deltId = blockDim.x / 2;
     }

   // в первом элементе массива хранится число частиц 
   nA = md->cells[iCA][0];
   nB = md->cells[iCB][0];
   nOwn = md->cells[iC][0]; // nA или nB

   // определяем диапазон частиц, которые будет считывать данный поток
  step = ceilf(2 * (double)nOwn / (double)blockDim.x);
   id0 = (threadIdx.x - deltId) * step;
   N = min(id0 + step, nOwn);
   for (i = id0; i < N; i++)
   {
      id = md->cells[iC][i+1];
      sIds[i] = id;		// сохраняем индексы частиц
      sXyz[i] = md->xyz[id];	// копируем координаты
      sFrs[i] = make_float3(0.0, 0.0, 0.0);    // обнуляем силы
   sTypes[i] = md->types[id];
   }
   __syncthreads();

   // взаимодействие между частицами внутри одной ячейки
   // потоки обрабатывают пары частиц начиная с индексов i=0 и j=1, с шагом по j равным числу потоков в блоке
   i = 0;
   j = threadIdx.x – deltId + 1;
   step1 = blockDim.x / 2;
   while (1)
   {
      while (j >= nOwn)
      {
          i++;
          if (i >= nOwn-1)
          {
              ex = 1;
              break;
          }
          j = i + 1 + j - nOwn;
      }
      if (ex) break;
      pair_in_cell(sXyz[i], @sFrs[i], sTypes[i], sXyz[j], @sFrs[j], sTypes[j], md, eVdW, eCoul);
      j = j + step1;
   }

   // вычисляем взаимодействия между частицами в разных ячейках
   step1 = ceilf(nA / blockDim.x);
   id1 = threadIdx.x * step1;
   N1 = min(id1 + step1, nA);
   for (i = id1; i < N1; i++)
     for (j = 0; j < nB; j++)
       pair_between_cell(sAxyz[i], &sAfrs[i], sAtypes[i], sBxyz[j], &sBfrs[j], sBtypes[j], md->cellShifts[blockId.x], md, eVdW, eCoul);

   __syncthreads();

   // копируем вычисленные силы в global memory
   for (i = id0; i < N; i++)
       md->frs[sIds[i]] = sFrs[i];

   // суммируем энергии внутри одного блока
   atomicAdd(&engCoul, eCoul);
   atomicAdd(&engVdw, eVdW);
   __syncthreads();

   // суммируем энергии блока в global memory
   if (!threadIdx.x) 
   {
      atomicAdd(&md->engCoul, engCoul);  
      atomicAdd(&md->engVdW, engVdW);
   }
}
 

Переменные вида float3PerCell – это произведение максимального числа частиц в ячейке на размер типа float3. Имеет смысл один раз вычислить их при инициализации системы. Функции pair_in_cell и pair_between_cell вычисляют взаимодействия между частицами в одной и разных ячейках, соответственно. В качестве параметров они получают координаты частиц, их типы и указатели куда записывать силы. Параметры для pair_between_cell включают также величины сдвигов, для учета периодических граничных условий. Коды для этих функций будут приведены чуть ниже. Для перебора пар частиц внутри ячейки я хотел пойти стандартным путем – разбить количество пар на равные интервалы и выдать каждому потоку по интервалу. Однако, чтобы получить индексы частиц i-j по номеру пары я пришел к квадратному уравнению и решил сделать по-другому: первый поток обрабатывает пару частиц с индексами 0-1, второй — 0-2 и т.д. Обработав пару i-j поток переходит к паре i-(j+N), где N – число потоков в блоке и если выходит за границу диапазона увеличивает индекс первой частицы в паре, см. схему:



Перебор пар между разными ячейками идёт следующим образом: все потоки получают свой интервал индексов частиц по ячейке A и рассчитывают взаимодействия между этими частицами и всеми частицами в ячейке B. Вторая функция вычисляет взаимодействия во всех оставшихся парах:

__global__ void cell_list2b(cudaMD* md)
{
    int iPair, iC, iCA, iCB, step, step1; 
    int i, j, id0, id1, N, N1, deltId;    
    int id;
    int nA, nB, nOwn; 

    __shared__ float engVdW, engCoul;
    float eVdW = 0.0;
    float eCoul = 0.0;

    extern __shared__ int shMem[];      // declaration of shared memory
    int* sAids = shMem;
    int* sAtypes = (int*)&sAids[md->maxAtPerCell];
    float3* sAxyz = (float3*)&sAtypes[md->maxAtPerCell];
    float3* sAfrs = (float3*)&sAxyz[md->maxAtPerCell];
    int* sBids = (int*)&sAfrs[md->maxAtPerCell];
    int* sBtypes = (int*)&sBids[md->maxAtPerCell];
    float3* sBxyz = (float3*)&sBtypes[md->maxAtPerCell];
    float3* sBfrs = (float3*)&sBxyz[md->maxAtPerCell];

    float3* sXyz;
    float3* sFrs;  
    int* sTypes;
    int* sIds;

    iPair = blockIdx.x + md->nPair1;
    if (threadIdx.x == 0)
    {
        engVdW = 0.0;
        engCoul = 0.0;
    }
    __syncthreads();

    iCA = md->cellPairs[iPair].x; 
    iCB = md->cellPairs[iPair].y;
    nA = md->cells[iCA][0];
    nB = md->cells[iCB][0];
    
    // нет смысла идти дальше, если хотя бы в одной ячейке нет частиц
    if (!nA)
       return;
    if (!nB)
        return;

    // сперва потоки копируют данные в shared_memory
    if (threadIdx.x < blockDim.x / 2)
    {
        iC = iCA;
        sXyz = sAxyz;
        sFrs = sAfrs;
        sTypes = sAtypes;
        sIds = sAids;
        deltId = 0;
    }
    else
    {
        iC = iCB;
        sXyz = sBxyz;
        sFrs = sBfrs;
        sTypes = sBtypes;
        sIds = sBids;
        deltId = blockDim.x / 2;
    }
    nOwn = md->cells[iC][0]; // nA or nB

    step = ceilf(2 * (double)nOwn / (double)blockDim.x);
    id0 = (threadIdx.x - deltId) * step;
    N = min(id0 + step, nOwn);
    for (i = id0; i < N; i++)
    {
        id = md->cells[iC][i + 1];
        sIds[i] = id;     
        sXyz[i] = md->xyz[id];    
        sFrs[i] = make_float3(0.0, 0.0, 0.0);    
        sTypes[i] = md->types[id];
    }
    __syncthreads();

    // вычисляем взаимодействия между атомами в разных ячейках
    step1 = ceilf((double)nA / (double)blockDim.x);
    id1 = threadIdx.x * step1;
    N1 = min(id1 + step1, nA);
    for (i = id1; i < N1; i++)
        for (j = 0; j < nB; j++)
        {
            pair_between_cell1(sAxyz[i], &sAfrs[i], sAtypes[i], sBxyz[j], &sBfrs[j], sBtypes[j], md->cellShifts[iPair], md, eVdW, eCoul);

        }
    __syncthreads();

   // копируем вычисленные силы в global memory
    for (i = id0; i < N; i++)
    {
        atomicAdd(&(md->frs[sIds[i]].x), sFrs[i].x);
        atomicAdd(&(md->frs[sIds[i]].y), sFrs[i].y);
        atomicAdd(&(md->frs[sIds[i]].z), sFrs[i].z);
    }

    atomicAdd(&engCoul, eCoul);
    atomicAdd(&engVdW, eVdW);
    __syncthreads();
 
    if (threadIdx.x == 0)
    {
        atomicAdd(&md->engCoul1, engCoul);
        atomicAdd(&md->engVdW, engVdW);
    }
}

Код почти аналогичен предыдущей функции, только теперь обрабатываются оставшиеся пары (индекс пары сдвигается на md->nPair1), вычисляется взаимодействие только между парой ячеек (внутри ячеек уже посчитано) и при копировании в глобальную память используем atomicAdd, чтобы суммировать результаты, а не перезаписать их. Теперь привожу оставшиеся функции:

__device__ void keep_in_cell(int index, float3 xyz, cudaMD* md)
// помещает index частицы с координатами xyz в cell list
{
    // определяем индекс ячейки
    int c = floor((double)xyz.x / (double)md->cSize.x) * md->cnYZ + floor((double)xyz.y / (double)md->cSize.y) * md->cNumber.z + floor((double)xyz.z / (double)md->cSize.z);

    int j = atomicAdd(&(md->cells[c][0]), 1);
    md->cells[c][j + 1] = index;
}

На основании координат вычисляем индекс той ячейки, куда попадает данная частица. Увеличиваем счетчик частиц в данной ячейке. Записываем индекс частицы в новую ячейку. сSize – размеры ячейки, cnYZ – произведение числа ячеек по измерениям Y и Z. Может быть стоит не вычислять индекс ячейки, а хранить их в виде 3-мерного массива в текстурной памяти, обращаясь по индексам вида [x / md->cSize.x]. Даст ли это прирост в производительности?

Функции, вычисляющие взаимодействия:

__device__ void pair_in_cell(float3 xyzA, float3* fxyzA, int typeA, float3 xyzB, float3* fxyzB, int typeB, cudaMD* md, float& engVdw, float& engCoul)
{
    float r, r2, f;
    float dx, dy, dz;
    cudaVdW* vdw = NULL;

    dx = xyzA.x - xyzB.x; 
    dy = xyzA.y - xyzB.y;
    dz = xyzA.z - xyzB.z;
    r2 = dx * dx + dy * dy + dz * dz;

    //if (r2 <= sim->r2Max)    
    //{
    r = 0.0;
    f = 0.0;

    // Кулоновское взаимодействие
    f += real_ewald(r2, r, md->chProd[typeA][typeB], md->alpha, engCoul);

    // Ван-дер-Ваальсово взаимодействие
    vdw = md->vdws[typeA][typeB];
    if (vdw != NULL)
        //if (r2 <= vdw->r2cut)   //! внутри ячейки эту проверку можно опустить
            f += vdw->feng_r(r2, r, vdw, engVdw);

    // в этом случае необходимо пользоваться atomic, поскольку одни частицы могут обрабатываться сразу разными потоками
    atomicAdd(&(fxyzA->x), f * dx);
    atomicAdd(&(fxyzB->x), -f * dx);
    atomicAdd(&(fxyzA->y), f * dy);
    atomicAdd(&(fxyzB->y), -f * dy);
    atomicAdd(&(fxyzA->z), f * dz);
    atomicAdd(&(fxyzB->z), -f * dz);
    //}
}

Вычисляем расстояние между частицами, подставляем его сначала в функцию для Кулоновского, а затем для Ван-дер-Ваальсова взаимодействия. Если размер ячейки выбрать таким образом, чтобы её диагональ не превышала дальности самого короткого типа взаимодействий, то проверки расстояний можно убрать (я специально оставил их закомментированными в коде). Замечу, что здесь я вычисляю квадрат расстояния (r2), а само расстояние (r) вычисляется в функции real_ewald. Массив chProd[i][j] хранит произведение зарядов соответствующих типов частиц. Возможно его тоже стоит поместить в текстурную память. Структура vdw хранит параметры парного потенциала и указатель на функцию (feng_r) для вычисления Ван-дер-Ваальсова взаимодействия.
Функция для вычисления взаимодействий между двумя разными ячейками:

__device__ void pair_between_cell1(float3 xyzA, float3* fxyzA, int typeA, float3 xyzB, float3* fxyzB, int typeB, float3 shift, cudaMD* md, float& engVdw, float& engCoul)
{
    float r, r2, f;
    float dx, dy, dz;
    cudaVdW* vdw;

    dx = xyzA.x - xyzB.x - shift.x;
    dy = xyzA.y - xyzB.y - shift.y;
    dz = xyzA.z - xyzB.z - shift.z;
    r2 = dx * dx + dy * dy + dz * dz;

    if (r2 <= md->r2Real)
    {
        r = 0.0;
        f = 0.0;

        f += real_ewald(r2, r, md->chProd[typeA][typeB], md->alpha, engCoul); 

        vdw = md->vdws[typeA][typeB];
        if (vdw != NULL)
            if (r2 <= vdw->r2cut)   
                f += vdw->feng_r(r2, r, vdw, engVdw);


        fxyzA->x += f * dx;
        atomicAdd(&(fxyzB->x), -f * dx);
        fxyzA->y += f * dy;
        atomicAdd(&(fxyzB->y), -f * dy);
        fxyzA->z += f * dz;
        atomicAdd(&(fxyzB->z), -f * dz);
    }
}

Код практически идентичен предыдущей функции, есть поправка к разнице координат для учета периодических граничных условий, раскомментированы проверки расстояний и, как мы помним из функций группы функций cell_list нет нужды применять atomicAdd к ячейке A, потому что частицы в ней обрабатываются независимо. Замечу, что в зависимости от расстояния между конкретными ячейками в паре можно опустить некоторые проверки расстояний и вычисления. Для этого можно написать слегка различные варианты этой функций. Привожу также функции вычисляющие силы и энергии парных потенциалов для потенциала Леннарда-Джонса и Букингема:

__device__ float cu_fe_lj(float r2, cudaVdW* vdw, float& eng)
// U = 4e[(s/r)^12 - (s/r)^6]
{
    float r2i = 1.0 / r2;
    float sr2 = vdw->p1 * r2i;
    float sr6 = sr2 * sr2 * sr2;

    eng += vdw->p0 * sr6 * (sr6 - 1.0);
    return vdw->p2 * r2i * sr6 * (2.0 * sr6 - 1.0);
}

__device__ float cu_fer_buck(float r2, float& r, cudaVdW* vdw, float& eng)
// U = A exp(-r/ro) - C/r^6
{
    float r2i = 1.0 / r2;
    float r4i = r2i * r2i;
    if (r == 0.0)    // если r неизвестно, вычислим
        r = sqrt(r2);

    eng += vdw->p0 * exp(-r / vdw->p1) - vdw->p2 * r4i * r2i;
    return vdw->p0 * exp(-r / vdw->p1) / r / vdw->p1 - 6.0 * vdw->p2 * r4i * r4i;
}

Функции возвращают силу (которую нужно домножить на проекции вектора расстояния) и добавляют энергию в переменную eng. Указатель на одну из этих функций хранится в структуре cudaVdW. Вообще говоря, для получения траекторий движения энергию можно и не вычислять, она нужна для разных статистик. Если мы уберем из этих функций выражения для энергии, расчеты, вероятно, ещё ускорятся. А можно сделать так: раз в несколько шагов вызывать полную функцию, а в остальные разы – только редуцированный аналог. Часть Кулоновского взаимодействия вычисляется функцией real_ewald:

__device__ float real_ewald(float r2, float& r, float chprd, float alpha, float& eng)
{
    float ar; 
    float erfcar;
    float kqq = chprd * d_Fcoul_scale; // q[i]*q[j]*1/4pie0;

    r = sqrt(r2);
    ar = alpha * r;
    erfcar = erfc(ar);

    eng += kqq * erfcar / r;
    return kqq / r / r2 * (erfcar + 2 * ar / d_sqrtpi * exp(-ar * ar));
}

Это отличается от знакомого закона Кулона, но об этом в отдельной статье.

Тестирование на реальном примере показало, что частиц в паре ячеек очень мало, чтобы загрузить все потоки, выделяемые на блок. Дальнейшая оптимизация функций cell_list2a и cell_list2b заключается в том, чтобы обрабатывать сразу несколько пар в одном блоке. Пока я написал только простейший вариант, где потоки просто делятся на несколько пар ячеек, т.е. один блок обрабатывает не одну пару, а скажем 4: 0-1, 1-2, 3-4, 5-6. Можно заметить, что в этом случае загрузив данные в shared memory можно заодно обработать пары 0-2, 0-3, 0-4 и т.д. Однако это требует дополнительных ухищрений на этапе подготовки данных, которые я ещё не продумал. Также если в функции cell_list2b загружать пары с общей ячейкой (например, те же 0-2, 0-3, 0-4 и т.д.), то можно сократить число обращений к глобальной памяти, скопировав данные из ячейки 0 только один раз.

Чтобы проверить всё ли правильно работает в моей реализации алгоритмов cell list я заодно написал функцию, вычисляющую взаимодействия наивным образом, просто перебирая все пары:

__global__ void all_pair(cudaMD* md)
{
    int i, j;
    int step = blockDim.x * gridDim.x;
    int ex = 0;  // exit flag
    int nAt = md->nAt;

    i = 0;
    j = blockIdx.x * blockDim.x + threadIdx.x + 1;
    while (1)
    {
        while (j >= nAt)
        {
            i++;
            if (i >= nAt - 1)
            {
                ex = 1;
                break;
            }
            j = i + 1 + j - nAt;
        }
        if (ex) break;
        global_pair(i, j, iStep, md);
        j = j + step;
    }
} 

__device__ void global_pair(int i, int j, cudaMD* md)
{
    float r, r2, f;
    float engCoul = 0;
    float engVdw = 0;
    cudaVdW* vdw;
    int ti, tj;

    float dx = md->xyz[i].x - md->xyz[j].x;
    float dy = md->xyz[i].y - md->xyz[j].y;
    float dz = md->xyz[i].z - md->xyz[j].z;

    // периодические граничные условия
    if (dx > md->halfLeng.x)
        dx -= md->leng.x;
    else
        if (dx < -md->halfLeng.x)
            dx += md->leng.x;

    if (dy > md->halfLeng.y)
        dy -= md->leng.y;
    else
        if (dy < -md->halfLeng.y)
            dy += md->leng.y;

    if (dz > md->halfLeng.z)
        dz -= md->leng.z;
    else
        if (dz < -md->halfLeng.z)
            dz += md->leng.z;

    r2 = dx * dx + dy * dy + dz * dz;

    if (r2 <= md->r2Real)    
    {
        r = 0.0;
        f = 0.0;

        ti = md->types[i];
        tj = md->types[j];

        f += real_ewald(r2, r, md->chProd[ti][tj], md->alpha, engCoul);

        vdw = md->vdws[ti][tj];
        if (vdw != NULL)
            if (r2 <= vdw->r2cut)   
                f += vdw->feng_r(r2, r, vdw, engVdw); 

        dx *= f;
        dy *= f;
        dz *= f;

        atomicAdd(&(md->frs[i].x), dx);
        atomicAdd(&(md->frs[j].x), -dx);
        atomicAdd(&(md->frs[i].y), dy);
        atomicAdd(&(md->frs[j].y), -dy);
        atomicAdd(&(md->frs[i].z), dz);
        atomicAdd(&(md->frs[j].z), -dz);

        atomicAdd(&(md->engCoul1), engCoul);
        atomicAdd(&(md->engVdW), engVdw);

    }
}

Результаты вычислений совпали. Правда, видимо, из-за разных последовательностей округлений координаты частиц в пределах 3-го знака после запятой могут варьироваться от запуска к запуску. Замечу, что на системе порядка 1400 частиц и радиуса обрезания 6 ангстрем функции cell_list2a и cell_list2b проигрывают по времени исполнения простому перебору примерно в два раза. При обработке в блоке сразу нескольких пар (8 и 16) скорости сравниваются и могут даже несколько опережать функцию all_pair. Однако это должно сильно зависеть от размеров системы. При увеличении числа атомов скорость алгоритмов cell list растет линейно, а простой перебор – квадратично. Кроме того, cell list очень эффективен при моделировании разряженных систем (большие расстояния между частицами), например, газов

Заключение

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

1. Суммирование по Эвальду.
2. Внутримолекулярное взаимодействие.
3. Разрыв и образование связей.
4. Электронный перенос.

Из литературы по МД могу порекомендовать книгу Daan Frenkel and Berend Smit. «Understanding Molecular Simulation: From Algorithms to Applications» и «Computer Simulation of Liquids» за авторством M.P. Allen, D. J. Tildesley. Для первой, насколько мне известно, существует перевод на русский.

Также выражаю благодарность Павлу Кудинову pavel_kudinov за его комментарий, призывающий «писать GPGPU код, простой и прямой как молоток», прочитав который, я, наконец, перешел к активным действиям, что и вылилось в эту статью.

UPD

27.06.2020. Благодаря комментарию avdx убрал обнуление глобальных параметров из функций и по-другому определил диапазон частиц выдаваемых потоку.