program matmul_gpu
  implicit none

  ! Three 16384 x 16384 FP32 matrices, ~3 GB total
  integer, parameter :: N = 16384

  real(4), allocatable :: A(:,:), B(:,:), C(:,:)

  ! k and acc must be declared here (in the host scope) so the LOCAL
  ! locality clause in DO CONCURRENT can reference them by name.
  integer   :: i, j, k
  real(4)   :: acc

  integer(8) :: t0, t1, rate
  real(8)    :: elapsed, gflops
  real(4)    :: checksum

  allocate(A(N,N), B(N,N), C(N,N))

  ! ------------------------------------------------------------------
  ! Initialization - with -stdpar=gpu, DO CONCURRENT runs on the GPU.
  ! Therefore, the compiler keeps the arrays in GPU-managed memory,
  ! and there is no host-to-device transfer before the multiply below.
  ! ------------------------------------------------------------------
  do concurrent (j = 1:N, i = 1:N)
    A(i,j) = real(mod((i-1) + (j-1),  97) + 1) * 0.01
    B(i,j) = real(mod((i-1) * (j-1), 101) + 1) * 0.01
    C(i,j) = 0.0
  end do

  call system_clock(t0, rate)

  ! ------------------------------------------------------------------
  ! Dense matrix multiply  C = A * B
  !
  ! LOCAL(k, acc) gives every concurrent iteration its own private copy
  ! of k and acc, which is required for correctness: without it, multiple
  ! threads would race on the same memory locations.
  !
  ! The compiler (nvfortran -stdpar=gpu) makes this into a CUDA kernel:
  !   - the (i, j) index space maps to the 2-D thread grid
  !   - the inner k-loop executes sequentially inside each GPU thread
  !
  ! On Blackwell (B200) with cc100, or Hopper (H100) with cc90, the
  ! ~268 million (i,j) pairs fill all FP32 CUDA cores many times over.
  ! ------------------------------------------------------------------
  do concurrent (j = 1:N, i = 1:N) local(k, acc)
    acc = 0.0
    do k = 1, N
      acc = acc + A(i,k) * B(k,j)
    end do
    C(i,j) = acc
  end do

  call system_clock(t1, rate)

  elapsed = real(t1 - t0, 8) / real(rate, 8)
  ! Each output element requires N multiply-adds = 2N floating-point ops.
  gflops  = 2.0d0 * real(N,8)**3 / (elapsed * 1.0d9)

  ! Spot-check a 16 x 16 corner to confirm the computation ran.
  checksum = 0.0
  do i = 1, 16
    do j = 1, 16
      checksum = checksum + C(i,j)
    end do
  end do

  write(*,'(a,i0,a,i0)') "Matrix size : ", N, " x ", N
  write(*,'(a,f8.3,a)')  "Elapsed     : ", elapsed, " s"
  write(*,'(a,f8.3,a)')  "Performance : ", gflops,  " GFLOP/s"
  write(*,'(a,es12.6)')  "Checksum    : ", checksum

  deallocate(A, B, C)
end program matmul_gpu