From 1d8c8560188eb603d0b741c612655e47b53a7a9f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= Date: Tue, 7 Mar 2023 15:45:53 -0700 Subject: [PATCH 1/5] ffn: Fuse some of the operations This provides about 4% speedup from 0.789 to 0.758s. --- gpt2.f90 | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/gpt2.f90 b/gpt2.f90 index 6b3f9b5..f7b0ffb 100644 --- a/gpt2.f90 +++ b/gpt2.f90 @@ -56,9 +56,16 @@ function linear(x, w, b) result(y) function ffn(x, fc_w, fc_b, proj_w, proj_b) result(y) real(sp), intent(in) :: x(:,:), fc_w(:,:), fc_b(:), proj_w(:,:), proj_b(:) real(sp) :: y(size(x,1),size(x,2)) -!real(sp) :: a(4*size(x,1),size(x,2)) -!a = gelu(linear(x, fc_w, fc_b)) -y = linear(gelu(linear(x, fc_w, fc_b)), proj_w, proj_b) +real(sp) :: a(4*size(x,1),size(x,2)) +integer :: i, j +!a = linear(x, fc_w, fc_b) +call matmul_2d(fc_w, x, a) +do j = 1, size(a,1) +do i = 1, size(a,2) + a(j,i) = gelu(a(j,i) + fc_b(j)) +end do +end do +y = linear(a, proj_w, proj_b) end function function attention(q, k, v, mask) result(y) From 2399607466236c02ea4c5fd827d1eccbd35dccb9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= Date: Tue, 7 Mar 2023 16:38:24 -0700 Subject: [PATCH 2/5] Implement faster tanh(x) function This provides about 20% speedup from 0.752s to 0.602s. --- gpt2.f90 | 24 +++++++++++++++++++++++- 1 file changed, 23 insertions(+), 1 deletion(-) diff --git a/gpt2.f90 b/gpt2.f90 index f7b0ffb..6ac016f 100644 --- a/gpt2.f90 +++ b/gpt2.f90 @@ -7,9 +7,31 @@ module gpt2_mod contains +elemental real(sp) function fast_tanh(x) result(y) +real(sp), intent(in) :: x +real(sp) :: x2 +if (x > 5) then + y = 1 +elseif (x < -5) then + y = -1 +else + x2 = x*x + y = x * (0.98569772605911309407 + x2 *(-0.2794500993392901382 & + + x2 * (6.8280504526399188164e-2 + x2 * (-1.0972014877337651823e-2 & + + x2 * (1.1132367134444316902e-3 + x2 * (-7.018851897305717565e-5 & + + x2 * (2.656616768082727089e-6 + x2 * (-5.5138381821615909058e-8 & + + x2 * 4.8162484477588665996e-10)))))))) + + ! Alternative implementation (division is expensive) + !a = x * (135135.0 + x2 * (17325.0 + x2 * (378.0 + x2))) + !b = 135135.0 + x2 * (62370.0 + x2 * (3150.0 + x2 * 28.0)) + !y = a / b +end if +end function + elemental real(sp) function gelu(x) result(y) real(sp), intent(in) :: x -y = 0.5_sp * x * (1 + tanh(sqrt(2 / pi) * (x + 0.044715_sp * x**3))) +y = 0.5_sp * x * (1 + fast_tanh(sqrt(2 / pi) * (x + 0.044715_sp * x**3))) end function function softmax(x) result(y) From 6e540ee0f843d30c0512f18ac0eb692c7bfb10db Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= Date: Tue, 7 Mar 2023 17:06:53 -0700 Subject: [PATCH 3/5] Implement fast_gelu --- gpt2.f90 | 26 +++++++++++++++++++++++++- 1 file changed, 25 insertions(+), 1 deletion(-) diff --git a/gpt2.f90 b/gpt2.f90 index 6ac016f..f9d2a1e 100644 --- a/gpt2.f90 +++ b/gpt2.f90 @@ -34,6 +34,30 @@ elemental real(sp) function gelu(x) result(y) y = 0.5_sp * x * (1 + fast_tanh(sqrt(2 / pi) * (x + 0.044715_sp * x**3))) end function +elemental real(sp) function fast_gelu(x) result(y) +real(sp), intent(in) :: x +real(sp) :: x2, x_ +x_ = sqrt(2 / pi) * (x + 0.044715_sp * x**3) +if (x_ > 5) then + y = x +elseif (x_ < -5) then + y = 0 +else + x2 = x_*x_ + y = x_ * (0.98569772605911309407 + x2 *(-0.2794500993392901382 & + + x2 * (6.8280504526399188164e-2 + x2 * (-1.0972014877337651823e-2 & + + x2 * (1.1132367134444316902e-3 + x2 * (-7.018851897305717565e-5 & + + x2 * (2.656616768082727089e-6 + x2 * (-5.5138381821615909058e-8 & + + x2 * 4.8162484477588665996e-10)))))))) + y = 0.5_sp * x * (1 + y) + + ! Alternative implementation (division is expensive) + !a = x * (135135.0 + x2 * (17325.0 + x2 * (378.0 + x2))) + !b = 135135.0 + x2 * (62370.0 + x2 * (3150.0 + x2 * 28.0)) + !y = a / b +end if +end function + function softmax(x) result(y) real(sp), intent(in) :: x(:,:) real(sp) :: y(size(x,1),size(x,2)) @@ -84,7 +108,7 @@ function ffn(x, fc_w, fc_b, proj_w, proj_b) result(y) call matmul_2d(fc_w, x, a) do j = 1, size(a,1) do i = 1, size(a,2) - a(j,i) = gelu(a(j,i) + fc_b(j)) + a(j,i) = fast_gelu(a(j,i) + fc_b(j)) end do end do y = linear(a, proj_w, proj_b) From 938963882450d33b0bac46daa47f93943bfd4c03 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= Date: Tue, 7 Mar 2023 17:12:35 -0700 Subject: [PATCH 4/5] Inline fast_gelu --- gpt2.f90 | 19 +++++++++++++++++-- 1 file changed, 17 insertions(+), 2 deletions(-) diff --git a/gpt2.f90 b/gpt2.f90 index f9d2a1e..00b2c00 100644 --- a/gpt2.f90 +++ b/gpt2.f90 @@ -102,13 +102,28 @@ function linear(x, w, b) result(y) function ffn(x, fc_w, fc_b, proj_w, proj_b) result(y) real(sp), intent(in) :: x(:,:), fc_w(:,:), fc_b(:), proj_w(:,:), proj_b(:) real(sp) :: y(size(x,1),size(x,2)) -real(sp) :: a(4*size(x,1),size(x,2)) +real(sp) :: a(4*size(x,1),size(x,2)), x0, x_, x2 integer :: i, j !a = linear(x, fc_w, fc_b) call matmul_2d(fc_w, x, a) do j = 1, size(a,1) do i = 1, size(a,2) - a(j,i) = fast_gelu(a(j,i) + fc_b(j)) + !a(j,i) = fast_gelu(a(j,i) + fc_b(j)) + x0 = a(j,i) + fc_b(j) + x_ = sqrt(2 / pi) * (x0 + 0.044715_sp * x0**3) + if (x_ > 5) then + a(j,i) = x0 + elseif (x_ < -5) then + a(j,i) = 0 + else + x2 = x_*x_ + a(j,i) = x_ * (0.98569772605911309407 + x2 *(-0.2794500993392901382 & + + x2 * (6.8280504526399188164e-2 + x2 * (-1.0972014877337651823e-2 & + + x2 * (1.1132367134444316902e-3 + x2 * (-7.018851897305717565e-5 & + + x2 * (2.656616768082727089e-6 + x2 * (-5.5138381821615909058e-8 & + + x2 * 4.8162484477588665996e-10)))))))) + a(j,i) = 0.5_sp * x0 * (1 + a(j,i)) + end if end do end do y = linear(a, proj_w, proj_b) From fdefc3e207d90eb9e0ef1ad344276981b2dc8189 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ond=C5=99ej=20=C4=8Cert=C3=ADk?= Date: Tue, 7 Mar 2023 17:14:03 -0700 Subject: [PATCH 5/5] Flip the loop order This gets from about 0.605s to 0.594s. --- gpt2.f90 | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gpt2.f90 b/gpt2.f90 index 00b2c00..ac25409 100644 --- a/gpt2.f90 +++ b/gpt2.f90 @@ -106,8 +106,8 @@ function ffn(x, fc_w, fc_b, proj_w, proj_b) result(y) integer :: i, j !a = linear(x, fc_w, fc_b) call matmul_2d(fc_w, x, a) -do j = 1, size(a,1) do i = 1, size(a,2) +do j = 1, size(a,1) !a(j,i) = fast_gelu(a(j,i) + fc_b(j)) x0 = a(j,i) + fc_b(j) x_ = sqrt(2 / pi) * (x0 + 0.044715_sp * x0**3)