Документация Engee

Неявный/явный решатель с ускорением CUDA для 2D-модели Билера — Рейтера

История вопроса

SciML — это набор оптимизированных библиотек Julia для решения обыкновенных дифференциальных уравнений (ОДУ). SciML предоставляет множество явных и неявных решателей, подходящих для решения задач ОДУ различных видов. Систему дифференциальных уравнений в частных производных (ДУЧП) можно привести к задаче ОДУ методом прямых (МП). Суть МП состоит в дискретизации пространственных производных (методом конечных разностей, конечных объемов или конечных элементов) для их преобразования в алгебраические уравнения с сохранением производных по времени без изменения. В получившихся дифференциальных уравнениях остается только одна независимая переменная (время), поэтому их можно решить с помощью решателя ODE. Статья Решение систем стохастических дифференциальных уравнений в частных производных и использование GPU в Julia содержит краткое введение в МП и использование GPU для ускорения решения ДУЧП в JuliaDiffEq. Здесь мы развернем эту тему, разработав неявный-явный (IMEX) решатель для двухмерной кардиоэлектрофизиологической модели и покажем, как использовать библиотеки CUDA для выполнения явной части модели на GPU.

Обратите внимание, что в этом руководстве не применяются методы IMEX более высоких порядков, встроенные в DifferentialEquations.jl. Вместо этого в нем демонстрируется, как вручную разделить уравнение в том случае, когда его явная часть имеет аналитическое (или приближенное) решение, что нередко встречается на практике.

Существуют сотни ионных моделей, описывающих электрическую активность сердца с различной степенью детализации. Большинство из них основаны на классической модели Ходжкина-Хаксли и представляют эволюцию различных переменных состояния во время в форме нелинейных ОДУ первого порядка. Вектор состояния для этих моделей включает в себя трансмембранный потенциал, управляющие переменные и концентрации ионов. Связь между клетками осуществляется только посредством трансмембранного потенциала и описывается в виде уравнения реакционной диффузии, представляющего собой параболическое ДУЧП:

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

где теперь является скалярной величиной. По своей природе эти модели работают с разными масштабами времени и поэтому относятся к категории жестких. Как правило, их решение находится с помощью явного метода Эйлера, обычно с замкнутой формой для интегрирования управляющих переменных (метод Раша-Ларсена, см. ниже). Решать такие задачи можно также с помощью полунеявных решателей ДУЧП (например, методом Кранка-Николсона в сочетании с итерационным решателем). Явные методы более высоких порядков, например метод Рунге-Кутты и линейный многошаговый метод, не позволяют преодолеть жесткость, и поэтому от них мало пользы.

В этом руководстве мы сначала разработаем решатель IMEX, задействующий только ЦП, а затем покажем, как перенести обработку явной части на GPU.

Модель Билера — Рейтера

В качестве примера мы выбрали вентрикулярную ионную модель Билера-Рейтера. Это классическая модель, впервые описанная в 1977 году и применяемая в качестве основы для множества других ионных моделей. В ней шесть переменных состояния, что делает ее достаточно интересной, но не слишком сложной для понимания основных аспектов задачи. Вот эти шесть переменных: трансмембранный потенциал ( ), ворота активации и инактивации натриевых каналов ( и , аналогично модели Ходжкина-Хаксли) с дополнительными воротами медленной инактивации ( ), ворота активации и инактивации кальциевых каналов ( и ), ворота зависимого от времени внутреннего выпрямляющего калиевого тока ( ) и внутриклеточная концентрация кальция ( ). Токов всего четыре: натриевый ( ), кальциевый ( ) и два калиевых: один зависимый от времени ( ), а другой фоновый независимый от времени ( ).

Решатель Билера — Рейтера только для ЦП

Начнем с разработки решателя IMEX, задействующего только ЦП. Основная идея заключается в том, что для обработки неявной части уравнения будет использоваться фреймворк DifferentialEquations, а явная часть будет обрабатываться отдельно с помощью кода для аналитической аппроксимации. Если аналитическая аппроксимация была бы неизвестна, можно было бы воспользоваться методами из этого списка.

Сначала определим константы модели.

const v0 = -84.624
const v1 = 10.0
const C_K1 = 1.0f0
const C_x1 = 1.0f0
const C_Na = 1.0f0
const C_s = 1.0f0
const D_Ca = 0.0f0
const D_Na = 0.0f0
const g_s = 0.09f0
const g_Na = 4.0f0
const g_NaC = 0.005f0
const ENa = 50.0f0 + D_Na
const γ = 0.5f0
const C_m = 1.0f0

Обратите внимание, что константы определяются как Float32, а не Float64. Причина в том, что в большинстве GPU гораздо больше ядер одинарной точности, чем двойной. Для обеспечения единообразной обработки с помощью ЦП и GPU мы также задаем большинство переменных состояний как Float32, кроме трансмембранного потенциала, который вычисляется с помощью неявного решателя из библиотеки Sundial и должен иметь тип Float64.

Структура состояния

Далее мы определяем структуру, которая будет содержать состояние. BeelerReuterCpu — это функтор, и мы определим функцию производной в качестве связанной функции.

mutable struct BeelerReuterCpu
    t::Float64              # время последнего временного шага для вычисления Δt
    diff_coef::Float64      # коэффициент диффузии (сила взаимодействия)

    C::Array{Float32, 2}    # внутриклеточная концентрация кальция
    M::Array{Float32, 2}    # ворота активации натриевого тока (m)
    H::Array{Float32, 2}    # ворота инактивации натриевого тока (h)
    J::Array{Float32, 2}    # ворота медленной инактивации натриевого тока (j)
    D::Array{Float32, 2}    # ворота активации кальциевого тока (d)
    F::Array{Float32, 2}    # ворота инактивации кальциевого тока (f)
    XI::Array{Float32, 2}   # внутренний выпрямляющий калиевый ток (iK1)

    Δu::Array{Float64, 2}   # заполнитель для лапласиана

    function BeelerReuterCpu(u0, diff_coef)
        self = new()

        ny, nx = size(u0)
        self.t = 0.0
        self.diff_coef = diff_coef

        self.C = fill(0.0001f0, (ny, nx))
        self.M = fill(0.01f0, (ny, nx))
        self.H = fill(0.988f0, (ny, nx))
        self.J = fill(0.975f0, (ny, nx))
        self.D = fill(0.003f0, (ny, nx))
        self.F = fill(0.994f0, (ny, nx))
        self.XI = fill(0.0001f0, (ny, nx))

        self.Δu = zeros(ny, nx)

        return self
    end
end

Лапласиан

Лапласиан конечной разности вычисляется на месте по пятиточечному шаблону. Применяется граничное условие Неймана. Обратите внимание, что для автоматизации этого шага можно было бы также воспользоваться DiffEqOperators.jl.

# пятиточечный шаблон
function laplacian(Δu, u)
    n1, n2 = size(u)

    # внутренние узлы
    for j in 2:(n2 - 1)
        for i in 2:(n1 - 1)
            @inbounds Δu[i, j] = u[i + 1, j] + u[i - 1, j] + u[i, j + 1] + u[i, j - 1] -
                                 4 * u[i, j]
        end
    end

    # левые и правые ребра
    for i in 2:(n1 - 1)
        @inbounds Δu[i, 1] = u[i + 1, 1] + u[i - 1, 1] + 2 * u[i, 2] - 4 * u[i, 1]
        @inbounds Δu[i, n2] = u[i + 1, n2] + u[i - 1, n2] + 2 * u[i, n2 - 1] - 4 * u[i, n2]
    end

    # верхние и нижние ребра
    for j in 2:(n2 - 1)
        @inbounds Δu[1, j] = u[1, j + 1] + u[1, j - 1] + 2 * u[2, j] - 4 * u[1, j]
        @inbounds Δu[n1, j] = u[n1, j + 1] + u[n1, j - 1] + 2 * u[n1 - 1, j] - 4 * u[n1, j]
    end

    # углы
    @inbounds Δu[1, 1] = 2 * (u[2, 1] + u[1, 2]) - 4 * u[1, 1]
    @inbounds Δu[n1, 1] = 2 * (u[n1 - 1, 1] + u[n1, 2]) - 4 * u[n1, 1]
    @inbounds Δu[1, n2] = 2 * (u[2, n2] + u[1, n2 - 1]) - 4 * u[1, n2]
    @inbounds Δu[n1, n2] = 2 * (u[n1 - 1, n2] + u[n1, n2 - 1]) - 4 * u[n1, n2]
end

Метод Раша — Ларсена

Мы используем явный решатель для всех переменных состояния, кроме трансмембранного потенциала, решение для которого находится с помощью неявного решателя. Явный решатель основан на специализированном экспоненциальном методе — методе Раша-Ларсена. Этот метод использует аппроксимацию модели для преобразования уравнения IMEX в форму, подходящую для неявного решателя ОДУ. Это сочетание неявного и явного методов образует специализированный решатель IMEX. Общие сведения об интегрировании IMEX см. в документации по решателям IMEX. Хотя для решения текущей задачи можно было бы воспользоваться общей моделью, подход с преобразованием в данном случае будет эффективнее и представляет практический интерес.

Метод Раша-Ларсена заменяет явное интегрирование методом Эйлера для управляющих переменных на прямое интегрирование. Отправной точкой является общее ОДУ для управляющих переменных в стиле ОДУ Ходжкина-Хаксли:

где  — универсальная управляющая переменная в диапазоне от 0 до 1, а и  — скорости реакций. Это уравнение можно записать в следующем виде:

где и  — это

и

Полагая, что и остаются постоянными на протяжении одного временного шага ( ), что является обоснованным допущением для большинства моделей сердечной активности, мы можем выполнить интегрирование напрямую:

Это и есть прием Раша-Ларсена. Обратите внимание: так как , это уравнение преобразуется в явную формулу Эйлера:

rush_larsen — это вспомогательная функция, которая использует метод Раша-Ларсена для интегрирования управляющих переменных.

@inline function rush_larsen(g, α, β, Δt)
    inf = α / (α + β)
    τ = 1.0f0 / (α + β)
    return clamp(g + (g - inf) * expm1(-Δt / τ), 0.0f0, 1.0f0)
end

Управляющие переменные обновляются, как указано ниже. и вычисляются согласно модели Билера-Рейтера, но конкретный способ их вычисления не представляет интереса в рамках данного руководства.

function update_M_cpu(g, v, Δt)
    # здесь требуется условие во избежание NaN при v == 47.0
    α = isapprox(v, 47.0f0) ? 10.0f0 : -(v + 47.0f0) / (exp(-0.1f0 * (v + 47.0f0)) - 1.0f0)
    β = (40.0f0 * exp(-0.056f0 * (v + 72.0f0)))
    return rush_larsen(g, α, β, Δt)
end

function update_H_cpu(g, v, Δt)
    α = 0.126f0 * exp(-0.25f0 * (v + 77.0f0))
    β = 1.7f0 / (exp(-0.082f0 * (v + 22.5f0)) + 1.0f0)
    return rush_larsen(g, α, β, Δt)
end

function update_J_cpu(g, v, Δt)
    α = (0.55f0 * exp(-0.25f0 * (v + 78.0f0))) / (exp(-0.2f0 * (v + 78.0f0)) + 1.0f0)
    β = 0.3f0 / (exp(-0.1f0 * (v + 32.0f0)) + 1.0f0)
    return rush_larsen(g, α, β, Δt)
end

function update_D_cpu(g, v, Δt)
    α = γ * (0.095f0 * exp(-0.01f0 * (v - 5.0f0))) / (exp(-0.072f0 * (v - 5.0f0)) + 1.0f0)
    β = γ * (0.07f0 * exp(-0.017f0 * (v + 44.0f0))) / (exp(0.05f0 * (v + 44.0f0)) + 1.0f0)
    return rush_larsen(g, α, β, Δt)
end

function update_F_cpu(g, v, Δt)
    α = γ * (0.012f0 * exp(-0.008f0 * (v + 28.0f0))) / (exp(0.15f0 * (v + 28.0f0)) + 1.0f0)
    β = γ * (0.0065f0 * exp(-0.02f0 * (v + 30.0f0))) / (exp(-0.2f0 * (v + 30.0f0)) + 1.0f0)
    return rush_larsen(g, α, β, Δt)
end

function update_XI_cpu(g, v, Δt)
    α = (0.0005f0 * exp(0.083f0 * (v + 50.0f0))) / (exp(0.057f0 * (v + 50.0f0)) + 1.0f0)
    β = (0.0013f0 * exp(-0.06f0 * (v + 20.0f0))) / (exp(-0.04f0 * (v + 20.0f0)) + 1.0f0)
    return rush_larsen(g, α, β, Δt)
end

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

function update_C_cpu(g, d, f, v, Δt)
    ECa = D_Ca - 82.3f0 - 13.0278f0 * log(g)
    kCa = C_s * g_s * d * f
    iCa = kCa * (v - ECa)
    inf = 1.0f-7 * (0.07f0 - g)
    τ = 1.0f0 / 0.07f0
    return g + (g - inf) * expm1(-Δt / τ)
end

Неявный решатель

Теперь пора определить функцию производной как функцию, связанную с BeelerReuterCpu. Для неявной части мы планируем использовать решатель CVODE_BDF. Как и другие итерационные методы, он многократно вызывает функцию производной с теми же значениями . Например, вот последовательность значений из показательного выполнения:

0,86830 0,86830 0,85485 0,85485 0,85485 0,86359 0,86359 0,86359 0,87233 0,87233 0,87233 0,88598 …​

Здесь каждый временной шаг вызывается три раза. Различаются два вида вызовов функции производной. При изменении управляющие переменные обновляются путем вызова update_gates_cpu:

function update_gates_cpu(u, XI, M, H, J, D, F, C, Δt)
    let Δt = Float32(Δt)
        n1, n2 = size(u)
        for j in 1:n2
            for i in 1:n1
                v = Float32(u[i, j])

                XI[i, j] = update_XI_cpu(XI[i, j], v, Δt)
                M[i, j] = update_M_cpu(M[i, j], v, Δt)
                H[i, j] = update_H_cpu(H[i, j], v, Δt)
                J[i, j] = update_J_cpu(J[i, j], v, Δt)
                D[i, j] = update_D_cpu(D[i, j], v, Δt)
                F[i, j] = update_F_cpu(F[i, j], v, Δt)

                C[i, j] = update_C_cpu(C[i, j], D[i, j], F[i, j], v, Δt)
            end
        end
    end
end

В свою очередь, переменная du обновляется на каждом временном шаге, так как она не зависит от .

# iK1 — это внутренний выпрямляющий калиевый ток
function calc_iK1(v)
    ea = exp(0.04f0 * (v + 85.0f0))
    eb = exp(0.08f0 * (v + 53.0f0))
    ec = exp(0.04f0 * (v + 53.0f0))
    ed = exp(-0.04f0 * (v + 23.0f0))
    return 0.35f0 * (4.0f0 * (ea - 1.0f0) / (eb + ec)
            +
            0.2f0 * (isapprox(v, -23.0f0) ? 25.0f0 : (v + 23.0f0) / (1.0f0 - ed)))
end

# ix1 — это независимый от времени фоновый калиевый ток
function calc_ix1(v, xi)
    ea = exp(0.04f0 * (v + 77.0f0))
    eb = exp(0.04f0 * (v + 35.0f0))
    return xi * 0.8f0 * (ea - 1.0f0) / eb
end

# iNa — это натриевый ток (как и в классической модели Ходжкина-Хаксли)
function calc_iNa(v, m, h, j)
    return C_Na * (g_Na * m^3 * h * j + g_NaC) * (v - ENa)
end

# iCa — это кальциевый ток
function calc_iCa(v, d, f, c)
    ECa = D_Ca - 82.3f0 - 13.0278f0 * log(c)    # ECa — это потенциал реверсии кальциевого тока
    return C_s * g_s * d * f * (v - ECa)
end

function update_du_cpu(du, u, XI, M, H, J, D, F, C)
    n1, n2 = size(u)

    for j in 1:n2
        for i in 1:n1
            v = Float32(u[i, j])

            # вычисляем отдельные токи
            iK1 = calc_iK1(v)
            ix1 = calc_ix1(v, XI[i, j])
            iNa = calc_iNa(v, M[i, j], H[i, j], J[i, j])
            iCa = calc_iCa(v, D[i, j], F[i, j], C[i, j])

            # общий ток
            I_sum = iK1 + ix1 + iNa + iCa

            # часть уравнения реакционной диффузии, представляющая реакцию
            du[i, j] = -I_sum / C_m
        end
    end
end

Наконец, мы помещаем все это в функцию производной, которая представляет собой вызов BeelerReuterCpu.

function (f::BeelerReuterCpu)(du, u, p, t)
    Δt = t - f.t

    if Δt != 0 || t == 0
        update_gates_cpu(u, f.XI, f.M, f.H, f.J, f.D, f.F, f.C, Δt)
        f.t = t
    end

    laplacian(f.Δu, u)

    # вычисляем часть, представляющую реакцию
    update_du_cpu(du, u, f.XI, f.M, f.H, f.J, f.D, f.F, f.C)

    # ...добавляем часть, представляющую диффузию
    du .+= f.diff_coef .* f.Δu
end

Результаты

Пора приступать к тестам! Нам необходимо определить начальный трансмембранный потенциал с помощью глобальных констант v0 и v1, которые представляют потенциалы в состоянии покоя и активации.

const N = 192;
u0 = fill(v0, (N, N));
u0[90:102, 90:102] .= v1;   # небольшой квадрат в центре области

Начальное условие представляет собой небольшой квадрат в центре области.

using Plots
heatmap(u0)

Далее определяем задачу:

using DifferentialEquations, Sundials

deriv_cpu = BeelerReuterCpu(u0, 1.0);
prob = ODEProblem(deriv_cpu, u0, (0.0, 50.0));

Для жестких уравнений реакционной диффузии отличным решателем является CVODE_BDF из библиотеки Sundial.

@time sol = solve(prob, CVODE_BDF(linear_solver = :GMRES), saveat = 100.0);
heatmap(sol.u[end])

Решатель Билера — Рейтера для ЦП и графического процессора

GPU отлично подходят для решения задач с высокой степенью параллелизма, но для моделей с высокой связностью они не так хороши. Мы собираемся неявную часть оставить на ЦП, а несвязный код для явной части выполнять на GPU с помощью библиотеки CUDAnative.

Графические процессоры и CUDA

В этом разделе дается краткий обзор принципов работы GPU (в частности, GPU от NVIDIA) и их программирования с помощью интерфейса CUDA в Julia. Читатели, знакомые с этими базовыми концепциями, могут пропустить этот раздел.

Начнем с рассмотрения аппаратной части типичного высококлассного GPU GTX 1080. У него четыре кластера обработки графики (эквивалентных отдельным ЦП), каждый из которых содержит пять потоковых мультипроцессоров (аналогичных ядрам ЦП). Каждый потоковый мультипроцессор имеет 128 ядер CUDA одинарной точности. Поэтому в сумме GTX 1080 имеет 4 x 5 x 128 = 2560 ядер CUDA. Максимальная теоретическая пропускная способность GTX 1080 составляет 8,87 терафлопса. Это значение рассчитывается для разогнанной тактовой частоты 1,733 МГц как 2 x 2560 x 1,733 МГц = 8,87 терафлопса. Коэффициент 2 добавлен по той причине, что за один такт могут выполняться две отдельные операции с плавающей запятой, умножение и сложение, в рамках одной операции FMA (fused-multiply-addition, умножение-сложение с однократным округлением). GTX 1080 также имеет 8192 МБ глобальной памяти, доступной всем ядрам (помимо локальной и общей памяти каждого потокового мультипроцессора).

Процесс работы приложения CUDA обычно выглядит так.

  1. В памяти ЦП определяются и инициализируются тензоры предметной области (многомерные массивы).

  2. Соответствующие тензоры размещаются в глобальной памяти GPU.

  3. Входные тензоры передаются из ЦП в соответствующие тензоры GPU.

  4. Вызываются ядра CUDA (то есть функции GPU, которые могут вызываться из ЦП), выполняющие операции с тензорами GPU.

  5. Итоговые тензоры передаются из GPU обратно в ЦП.

  6. Тензоры обрабатываются в ЦП.

  7. При необходимости повторяются шаги 3—​6.

Некоторые библиотеки, такие как ArrayFire, обеспечивают высокоуровневую абстракцию, скрывающую шаги 2—​5. Однако в данном случае мы выберем низкоуровневый подход. Используя CUDA, мы получаем более детальный контроль и более высокую производительность. Обратной стороной является необходимость реализовывать каждый шаг вручную.

CuArray — это тонкий слой абстракции поверх API CUDA, который позволяет определять тензоры на стороне GPU и копировать данных в них и из них, но не реализует операции с тензорами. CUDAnative — это компилятор, который транслирует функции Julia, выступающие в роли ядер CUDA, в ptx (высокоуровневый ассемблерный язык CUDA).

Код CUDA

Залогом быстрого выполнения программ CUDA является минимизация операций передачи данных между памятью ЦП и GPU и обращений к глобальной памяти. Неявный решатель в настоящее время задействует только ЦП, но ему требуется доступ только к трансмембранному потенциалу. Остальные переменные состояния размещаются в памяти GPU.

Мы преобразуем BeelerReuterCpu в BeelerReuterGpu, определяя переменные состояния как CuArray вместо стандартных массивов Julia. Для ясности перед именем каждой переменной, определяемой в GPU, добавляется префикс d_. Обратите внимание, что служит для временного хранения лапласиана и остается на стороне ЦП.

using CUDA

mutable struct BeelerReuterGpu <: Function
    t::Float64                  # время последнего временного шага для вычисления Δt
    diff_coef::Float64          # коэффициент диффузии (сила взаимодействия)

    d_C::CuArray{Float32, 2}    # внутриклеточная концентрация кальция
    d_M::CuArray{Float32, 2}    # ворота активации натриевого тока (m)
    d_H::CuArray{Float32, 2}    # ворота инактивации натриевого тока (h)
    d_J::CuArray{Float32, 2}    # ворота медленной инактивации натриевого тока (j)
    d_D::CuArray{Float32, 2}    # ворота активации кальциевого тока (d)
    d_F::CuArray{Float32, 2}    # ворота инактивации кальциевого тока (f)
    d_XI::CuArray{Float32, 2}   # внутренний выпрямляющий калиевый ток (iK1)

    d_u::CuArray{Float64, 2}    # заполнитель для u в памяти устройства
    d_du::CuArray{Float64, 2}   # заполнитель для d_u в памяти устройства

    Δv::Array{Float64, 2}       # заполнитель для градиента напряжения

    function BeelerReuterGpu(u0, diff_coef)
        self = new()

        ny, nx = size(u0)
        @assert (nx % 16 == 0) && (ny % 16 == 0)
        self.t = 0.0
        self.diff_coef = diff_coef

        self.d_C = CuArray(fill(0.0001f0, (ny, nx)))
        self.d_M = CuArray(fill(0.01f0, (ny, nx)))
        self.d_H = CuArray(fill(0.988f0, (ny, nx)))
        self.d_J = CuArray(fill(0.975f0, (ny, nx)))
        self.d_D = CuArray(fill(0.003f0, (ny, nx)))
        self.d_F = CuArray(fill(0.994f0, (ny, nx)))
        self.d_XI = CuArray(fill(0.0001f0, (ny, nx)))

        self.d_u = CuArray(u0)
        self.d_du = CuArray(zeros(ny, nx))

        self.Δv = zeros(ny, nx)

        return self
    end
end

Функция лапласиана остается без изменений. Главным изменением в явных решателях для управляющих переменных является то, что перед функциями exp и expm1 добавляется префикс CUDAnative.. Это техническая недоработка, которая, надеемся, будет исправлена в будущем.

function rush_larsen_gpu(g, α, β, Δt)
    inf = α / (α + β)
    τ = 1.0 / (α + β)
    return clamp(g + (g - inf) * CUDAnative.expm1(-Δt / τ), 0.0f0, 1.0f0)
end

function update_M_gpu(g, v, Δt)
    # здесь требуется условие во избежание NaN при v == 47.0
    α = isapprox(v, 47.0f0) ? 10.0f0 :
        -(v + 47.0f0) / (CUDAnative.exp(-0.1f0 * (v + 47.0f0)) - 1.0f0)
    β = (40.0f0 * CUDAnative.exp(-0.056f0 * (v + 72.0f0)))
    return rush_larsen_gpu(g, α, β, Δt)
end

function update_H_gpu(g, v, Δt)
    α = 0.126f0 * CUDAnative.exp(-0.25f0 * (v + 77.0f0))
    β = 1.7f0 / (CUDAnative.exp(-0.082f0 * (v + 22.5f0)) + 1.0f0)
    return rush_larsen_gpu(g, α, β, Δt)
end

function update_J_gpu(g, v, Δt)
    α = (0.55f0 * CUDAnative.exp(-0.25f0 * (v + 78.0f0))) /
        (CUDAnative.exp(-0.2f0 * (v + 78.0f0)) + 1.0f0)
    β = 0.3f0 / (CUDAnative.exp(-0.1f0 * (v + 32.0f0)) + 1.0f0)
    return rush_larsen_gpu(g, α, β, Δt)
end

function update_D_gpu(g, v, Δt)
    α = γ * (0.095f0 * CUDAnative.exp(-0.01f0 * (v - 5.0f0))) /
        (CUDAnative.exp(-0.072f0 * (v - 5.0f0)) + 1.0f0)
    β = γ * (0.07f0 * CUDAnative.exp(-0.017f0 * (v + 44.0f0))) /
        (CUDAnative.exp(0.05f0 * (v + 44.0f0)) + 1.0f0)
    return rush_larsen_gpu(g, α, β, Δt)
end

function update_F_gpu(g, v, Δt)
    α = γ * (0.012f0 * CUDAnative.exp(-0.008f0 * (v + 28.0f0))) /
        (CUDAnative.exp(0.15f0 * (v + 28.0f0)) + 1.0f0)
    β = γ * (0.0065f0 * CUDAnative.exp(-0.02f0 * (v + 30.0f0))) /
        (CUDAnative.exp(-0.2f0 * (v + 30.0f0)) + 1.0f0)
    return rush_larsen_gpu(g, α, β, Δt)
end

function update_XI_gpu(g, v, Δt)
    α = (0.0005f0 * CUDAnative.exp(0.083f0 * (v + 50.0f0))) /
        (CUDAnative.exp(0.057f0 * (v + 50.0f0)) + 1.0f0)
    β = (0.0013f0 * CUDAnative.exp(-0.06f0 * (v + 20.0f0))) /
        (CUDAnative.exp(-0.04f0 * (v + 20.0f0)) + 1.0f0)
    return rush_larsen_gpu(g, α, β, Δt)
end

function update_C_gpu(c, d, f, v, Δt)
    ECa = D_Ca - 82.3f0 - 13.0278f0 * CUDAnative.log(c)
    kCa = C_s * g_s * d * f
    iCa = kCa * (v - ECa)
    inf = 1.0f-7 * (0.07f0 - c)
    τ = 1.0f0 / 0.07f0
    return c + (c - inf) * CUDAnative.expm1(-Δt / τ)
end

Аналогичным образом изменяются функции для вычисления отдельных токов путем добавления префикса CUDAnative.

# iK1 — это внутренний выпрямляющий калиевый ток
function calc_iK1(v)
    ea = CUDAnative.exp(0.04f0 * (v + 85.0f0))
    eb = CUDAnative.exp(0.08f0 * (v + 53.0f0))
    ec = CUDAnative.exp(0.04f0 * (v + 53.0f0))
    ed = CUDAnative.exp(-0.04f0 * (v + 23.0f0))
    return 0.35f0 * (4.0f0 * (ea - 1.0f0) / (eb + ec)
            +
            0.2f0 * (isapprox(v, -23.0f0) ? 25.0f0 : (v + 23.0f0) / (1.0f0 - ed)))
end

# ix1 — это независимый от времени фоновый калиевый ток
function calc_ix1(v, xi)
    ea = CUDAnative.exp(0.04f0 * (v + 77.0f0))
    eb = CUDAnative.exp(0.04f0 * (v + 35.0f0))
    return xi * 0.8f0 * (ea - 1.0f0) / eb
end

# iNa — это натриевый ток (как и в классической модели Ходжкина-Хаксли)
function calc_iNa(v, m, h, j)
    return C_Na * (g_Na * m^3 * h * j + g_NaC) * (v - ENa)
end

# iCa — это кальциевый ток
function calc_iCa(v, d, f, c)
    ECa = D_Ca - 82.3f0 - 13.0278f0 * CUDAnative.log(c)    # ECa — это потенциал реверсии кальциевого тока
    return C_s * g_s * d * f * (v - ECa)
end

Ядра CUDA

Программа CUDA не работает напрямую с кластерами обработки графики и потоковыми мультипроцессорами. Логически программа CUDA представляется в виде блоков и потоков. При запуске ядра CUDA необходимо указать количество блоков и потоков. Каждый поток выполняется на одном ядре CUDA. Потоки логически объединяются в блоки, которые, в свою очередь, образуют сетку. Сетка полностью представляет требуемую область.

Каждый поток может определить свою логическую координату с помощью нескольких предварительно определенных переменных индексирования (threadIdx, blockIdx, blockDim и gridDim) в C/C++ и соответствующих функций (например, threadIdx()) в Julia. Эти переменные и функции определяются автоматически для каждого потока и могут возвращать разные значения в зависимости от вызывающего потока. Возвращаемое значение этих функций представляет собой одно-, двух- или трехмерную структуру, к элементам которой можно обращаться посредством .x, .y и .z (в случае с одномерной структурой .x возвращает фактический индекс, а .y и .z просто возвращают 1). Например, если ядро развернуто в 128 блоках с 256 потоками на блок, каждому потоку будут доступны следующие значения:

    gridDim.x = 128;
    blockDim=256;

blockIdx.x будет иметь диапазон от 0 до 127 в C/C++ и от 1 до 128 в Julia. Аналогичным образом, threadIdx.x будет иметь диапазон от 0 до 255 в C/C++ и от 1 до 256 в Julia.

Поток C/C++ может вычислить свой индекс следующим образом:

    int idx = blockDim.x * blockIdx.x + threadIdx.x;

В Julia следует учитывать, что отсчет начинается от 1. Поэтому применяется следующая формула.

    idx = (blockIdx().x-UInt32(1)) * blockDim().x + threadIdx().x

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

В версии решателя для GPU каждый поток работает с отдельным элементом среды, который индексируется по паре значений (x, y). Функции update_gates_gpu и update_du_gpu очень похожи на свои аналоги для ЦП, но по сути представляют собой ядра CUDA, в которых циклы for заменяются индексированием CUDA. Обратите внимание, что ядра CUDA не могут возвращать значения, поэтому в конце стоит nothing.

function update_gates_gpu(u, XI, M, H, J, D, F, C, Δt)
    i = (blockIdx().x - UInt32(1)) * blockDim().x + threadIdx().x
    j = (blockIdx().y - UInt32(1)) * blockDim().y + threadIdx().y

    v = Float32(u[i, j])

    let Δt = Float32(Δt)
        XI[i, j] = update_XI_gpu(XI[i, j], v, Δt)
        M[i, j] = update_M_gpu(M[i, j], v, Δt)
        H[i, j] = update_H_gpu(H[i, j], v, Δt)
        J[i, j] = update_J_gpu(J[i, j], v, Δt)
        D[i, j] = update_D_gpu(D[i, j], v, Δt)
        F[i, j] = update_F_gpu(F[i, j], v, Δt)

        C[i, j] = update_C_gpu(C[i, j], D[i, j], F[i, j], v, Δt)
    end
    nothing
end

function update_du_gpu(du, u, XI, M, H, J, D, F, C)
    i = (blockIdx().x - UInt32(1)) * blockDim().x + threadIdx().x
    j = (blockIdx().y - UInt32(1)) * blockDim().y + threadIdx().y

    v = Float32(u[i, j])

    # вычисляем отдельные токи
    iK1 = calc_iK1(v)
    ix1 = calc_ix1(v, XI[i, j])
    iNa = calc_iNa(v, M[i, j], H[i, j], J[i, j])
    iCa = calc_iCa(v, D[i, j], F[i, j], C[i, j])

    # общий ток
    I_sum = iK1 + ix1 + iNa + iCa

    # часть уравнения реакционной диффузии, представляющая реакцию
    du[i, j] = -I_sum / C_m
    nothing
end

Неявный решатель

Наконец, мы изменяем функцию производной для копирования переменной u в GPU, копирования переменной du обратно и вызова ядер CUDA.

function (f::BeelerReuterGpu)(du, u, p, t)
    L = 16   # размер блока
    Δt = t - f.t
    copyto!(f.d_u, u)
    ny, nx = size(u)

    if Δt != 0 || t == 0
        @cuda blocks=(ny ÷ L, nx ÷ L) threads=(L, L) update_gates_gpu(f.d_u, f.d_XI, f.d_M,
                                                                      f.d_H, f.d_J, f.d_D,
                                                                      f.d_F, f.d_C, Δt)
        f.t = t
    end

    laplacian(f.Δv, u)

    # вычисляем часть, представляющую реакцию
    @cuda blocks=(ny ÷ L, nx ÷ L) threads=(L, L) update_du_gpu(f.d_du, f.d_u, f.d_XI, f.d_M,
                                                               f.d_H, f.d_J, f.d_D, f.d_F,
                                                               f.d_C)

    copyto!(du, f.d_du)

    # ...добавляем часть, представляющую диффузию
    du .+= f.diff_coef .* f.Δv
end

Все готово для тестирования!

using DifferentialEquations, Sundials

deriv_gpu = BeelerReuterGpu(u0, 1.0);
prob = ODEProblem(deriv_gpu, u0, (0.0, 50.0));
@time sol = solve(prob, CVODE_BDF(linear_solver = :GMRES), saveat = 100.0);
heatmap(sol.u[end])

Заключение

Благодаря выполнению явной части решателя IMEX в GPU достигается шестикратное повышение скорости. Главным узким местом при таком подходе является взаимодействие между ЦП и GPU. В текущей форме не все внутренние механизмы метода используют ускорение GPU. В частности, неявные уравнения, решаемые GMRES, обрабатываются ЦП. Такая комбинированная обработка также увеличивает объем данных, передаваемых между GPU и ЦП (при каждом вызове f). Компиляция всего решателя ОДУ для GPU могла бы решить обе эти проблемы и, возможно, еще более существенно повысить скорость выполнения. Разработчики JuliaDiffEq в настоящее время работают над решениями для смягчения этих проблем, но они будут совместимы только с собственными решателями Julia (но не с решателями Sundials).