Skip to content

Re-factor of MOM_ANN #3

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Pperezhogin merged 6 commits into from
Jun 16, 2025
Merged

Re-factor of MOM_ANN #3

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
5b8212a
Moved MOM_ANN.F90 to src/framework/
adcroft May 2, 2025
f87c65a
Minor refactor of MOM_ANN
adcroft Apr 29, 2025
c4984c4
Adds unit tests and timing test to MOM_ANN
adcroft Apr 30, 2025
ea1cc31
Adding multiple forms of inference
adcroft May 2, 2025
4e98437
Renamed ANN variants and added some module documentation
adcroft May 9, 2025
da17218
Removed alternative variants of ANN in favor of optimized
adcroft Jun 2, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
46 changes: 23 additions & 23 deletions config_src/drivers/timing_tests/time_MOM_ANN.F90
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,9 +4,9 @@ program time_MOM_ANN

use MOM_ANN, only : ANN_CS
use MOM_ANN, only : ANN_allocate, ANN_apply, ANN_end
use MOM_ANN, only : ANN_apply_vector_v1, ANN_apply_vector_v2, ANN_apply_vector_v3
use MOM_ANN, only : ANN_apply_vector_v4
use MOM_ANN, only : ANN_apply_array_v1, ANN_apply_array_v2, ANN_apply_array_v3
use MOM_ANN, only : ANN_apply_vector_orig, ANN_apply_vector_oi, ANN_apply_vector_io
use MOM_ANN, only : ANN_apply_vector_oit
use MOM_ANN, only : ANN_apply_array_soi, ANN_apply_array_sio, ANN_apply_array_ois
use MOM_ANN, only : ANN_random

implicit none
Expand Down Expand Up @@ -57,35 +57,35 @@ program time_MOM_ANN
! Fastest variants on Intel Xeon W-2223 CPU @ 3.60GHz (gfortran-13.2 -O3)
! | vector(nxy=1) | nxy = 10 | nxy = 100
! ----------------------------------------------------------------------------
! Small ANN | vector_v2 | array_v1 | array_v2
! Shallow-wide ANN | vector_v2 | array_v3 | array_v2
! Deep ANN | vector_v2 | array_v3 | array_v2
! Small ANN | vector_oi | array_soi | array_sio
! Shallow-wide ANN | vector_oi | array_ois | array_sio
! Deep ANN | vector_oi | array_ois | array_sio

write(*,'(a)') "{"

call time_ANN(nlayers, nin, layer_width, nout, nsamp, nits, nxy, &
0, "MOM_ANN:ANN_apply(vector)")
write(*,"(',')")
call time_ANN(nlayers, nin, layer_width, nout, nsamp, nits, nxy, &
1, "MOM_ANN:ANN_apply_vector_v1(array)")
1, "MOM_ANN:ANN_apply_vector_orig(array)")
write(*,"(',')")
call time_ANN(nlayers, nin, layer_width, nout, nsamp, nits, nxy, &
2, "MOM_ANN:ANN_apply_vector_v2(array)")
2, "MOM_ANN:ANN_apply_vector_oi(array)")
write(*,"(',')")
call time_ANN(nlayers, nin, layer_width, nout, nsamp, nits, nxy, &
3, "MOM_ANN:ANN_apply_vector_v3(array)")
3, "MOM_ANN:ANN_apply_vector_io(array)")
write(*,"(',')")
call time_ANN(nlayers, nin, layer_width, nout, nsamp, nits, nxy, &
4, "MOM_ANN:ANN_apply_vector_v4(array)")
4, "MOM_ANN:ANN_apply_vector_oit(array)")
write(*,"(',')")
call time_ANN(nlayers, nin, layer_width, nout, nsamp, nits, nxy, &
11, "MOM_ANN:ANN_apply_array_v1(array)")
11, "MOM_ANN:ANN_apply_array_soi(array)")
write(*,"(',')")
call time_ANN(nlayers, nin, layer_width, nout, nsamp, nits, nxy, &
12, "MOM_ANN:ANN_apply_array_v2(array)")
12, "MOM_ANN:ANN_apply_array_sio(array)")
write(*,"(',')")
call time_ANN(nlayers, nin, layer_width, nout, nsamp, nits, nxy, &
13, "MOM_ANN:ANN_apply_array_v3(array)")
13, "MOM_ANN:ANN_apply_array_ois(array)")
write(*,"()")

write(*,'(a)') "}"
Expand Down Expand Up @@ -123,7 +123,7 @@ subroutine time_ANN(nlayers, nin, width, nout, nsamp, nits, nxy, impl, label)
real :: tmin, tmax, tmean, tstd ! Min, max, mean, and standard deviation, of CPU times [s]
integer :: asamp ! Actual samples of timings
integer :: aits ! Actual iterations
real :: flops ! Operations per sec estimated from parameters [# s-1]
real :: words_per_sec ! Operations per sec estimated from parameters [# s-1]

widths(:) = width
widths(1) = nin
Expand Down Expand Up @@ -154,52 +154,52 @@ subroutine time_ANN(nlayers, nin, width, nout, nsamp, nits, nxy, impl, label)
call cpu_time(start)
do iter = 1, aits ! Make many passes to reduce sampling error
do ij = 1, nxy
call ANN_apply_vector_v1(x_fs(:,ij), y_fs(:,ij), ANN)
call ANN_apply_vector_orig(x_fs(:,ij), y_fs(:,ij), ANN)
enddo
enddo
call cpu_time(finish)
case (2)
call cpu_time(start)
do iter = 1, aits ! Make many passes to reduce sampling error
do ij = 1, nxy
call ANN_apply_vector_v2(x_fs(:,ij), y_fs(:,ij), ANN)
call ANN_apply_vector_oi(x_fs(:,ij), y_fs(:,ij), ANN)
enddo
enddo
call cpu_time(finish)
case (3)
call cpu_time(start)
do iter = 1, aits ! Make many passes to reduce sampling error
do ij = 1, nxy
call ANN_apply_vector_v3(x_fs(:,ij), y_fs(:,ij), ANN)
call ANN_apply_vector_io(x_fs(:,ij), y_fs(:,ij), ANN)
enddo
enddo
call cpu_time(finish)
case (4)
call cpu_time(start)
do iter = 1, aits ! Make many passes to reduce sampling error
do ij = 1, nxy
call ANN_apply_vector_v4(x_fs(:,ij), y_fs(:,ij), ANN)
call ANN_apply_vector_oit(x_fs(:,ij), y_fs(:,ij), ANN)
enddo
enddo
call cpu_time(finish)
case (11)
call cpu_time(start)
do iter = 1, aits ! Make many passes to reduce sampling error
call ANN_apply_array_v1(nxy, x_sf(:,:), y_sf(:,:), ANN)
call ANN_apply_array_soi(nxy, x_sf(:,:), y_sf(:,:), ANN)
enddo
call cpu_time(finish)
asamp = nsamp * aits ! Account for working on whole arrays
case (12)
call cpu_time(start)
do iter = 1, aits ! Make many passes to reduce sampling error
call ANN_apply_array_v2(nxy, x_sf(:,:), y_sf(:,:), ANN)
call ANN_apply_array_sio(nxy, x_sf(:,:), y_sf(:,:), ANN)
enddo
call cpu_time(finish)
asamp = nsamp * aits ! Account for working on whole arrays
case (13)
call cpu_time(start)
do iter = 1, aits ! Make many passes to reduce sampling error
call ANN_apply_array_v3(nxy, x_fs(:,:), y_fs(:,:), ANN)
call ANN_apply_array_ois(nxy, x_fs(:,:), y_fs(:,:), ANN)
enddo
call cpu_time(finish)
asamp = nsamp * aits ! Account for working on whole arrays
Expand All @@ -217,15 +217,15 @@ subroutine time_ANN(nlayers, nin, width, nout, nsamp, nits, nxy, impl, label)
tstd = tstd / real(nsamp) ! convert to mean of squares
tstd = tstd - tmean**2 ! convert to variance
tstd = sqrt( tstd * real(nsamp) / real(nsamp-1) ) ! convert to standard deviation
flops = ANN%parameters / tmean
words_per_sec = ANN%parameters / ( tmean * 1024 * 1024 )

@Pperezhogin Pperezhogin Jun 9, 2025 •
edited
Loading

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Seems that this metric changed its meaning. Before it was measuring number of operators per second (flops) and now it is the memory throughput. I am not sure which metric would be the most relevant as it is unclear a priori if this code will be compute-bounded or memory-bounded. Is this words_per_sec metric supposed to be compared to L1/L2/L3 cache throughput?

I somewhat find Gflops to be more relevant metric as its range on single CPU core is clearly defined from approximately 0.5Gflops for scalar operators (typically, in scalar code per one floating point operation there are ~5 service operations, matmul benchmark) up to approximately 50-100Gflops for FMA instructions in longest vector registers assuming no memory transfer between registers and cache. Typical ocean model has 3Gflops performance on average which is much better than scalar code but still far away from the compute bound. I would say a metric of success for ANN module is to be more efficient than ocean model on average, i.e. be in a range of
3Gflops-100Gflops

May be, if we want Gflops, we may need to estimate the number of floating point operations, which is for matmul not the number of parameters in matrix, but approximately twice (one add and one multiply per element).

@adcroft adcroft Jun 10, 2025

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As @marshallward pointed out, what I'm calculating is the number of words for storage. When reporting Gflops, I had made the assumption that this was likely the number of multiply-adds, but as you say there is ambiguity in whether an FMA should count as one or two ops. Switching to memory processed avoided the ambiguity but in truth, it's probably better to just return the times (as we do in the other tests).


write(*,"(2x,3a)") '"', trim(label), '": {'
write(*,"(4x,a,1pe11.4,',')") '"min": ', tmin
write(*,"(4x,a,1pe11.4,',')") '"mean":', tmean
write(*,"(4x,a,1pe11.4,',')") '"std": ', tstd
write(*,"(4x,a,i0,',')") '"n_samples": ', asamp
write(*,"(4x,a,1pe11.4,',')") '"max": ', tmax
write(*,"(4x,a,1pe11.4,'}')", advance="no") '"flops": ', flops
write(*,"(4x,a,1pe11.4,'}')", advance="no") '"MBps": ', words_per_sec

end subroutine time_ANN

Expand Down
Loading