Python-CUDA-Algebraic-Multigrid/Example_AMG_PoissonProblem.py at main · KevinGFripp/Python-CUDA-Algebraic-Multigrid · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
from Laplacians.Laplacian2D import laplacian_2d
import numpy as np
import matplotlib.pyplot as plt
from Multigrid.HierarchyOfGrids import hierarchy_of_grids
from BICGSTAB_L.AMG_BICGSTAB_L import amg_bicgstab_l as amg_bicgstab_cpu
from BICGSTAB_L.AMG_BICGSTAB_L_GPU import amg_bicgstab_l_gpu as amg_bicgstab
from Multigrid.GridsToGPU import grids_to_gpu
from cupy import asarray,zeros_like
import time as timer


if __name__ == '__main__':
    ## Solve del2 * x = b on a regular grid
    ## Comparison between CPU performance and GPU performance

    #config
    runs = 100
    tol = 1e-5

    # grid size
    Nx = 512
    Ny = 512

    ##setup
    # make 2D laplacian matrix
    del2 = laplacian_2d(Nx, Ny)

    x0 = np.zeros((Nx * Ny), dtype=np.float32)
    solution_cpu = np.zeros_like(x0)

    # rhs with 1s on top boundary,-1s on bottom boundary, partial left-right boundaries of -1.0
    b = np.zeros((Nx * Ny), dtype=np.float32)
    b[0:Ny] = -1.0
    b[(Nx-1)*Ny:(Nx-1)*Ny + Ny] = -1.0
    b[(Nx//4)*Ny:(3*Nx//4)*Ny:Ny] = 1.0
    b[(Nx // 4) * Ny +Ny-1:(3 * Nx // 4) * Ny +Ny-1:Ny] = 1.0


    tic = timer.time()

    multigrid, numgrids = hierarchy_of_grids(del2,b,x0,Nx,Ny)
    toc = timer.time()
    print('Setup time = ',toc-tic,' s')


    print('--Multigrid levels--')
    [print(multigrid[n].Nx,'x',multigrid[n].Ny) for n in range(numgrids)]
    print('--------------------')

    # move to GPU
    x0_gpu = asarray(x0, dtype=np.float32)
    b_gpu = asarray(b, dtype=np.float32)
    multigrid_gpu = grids_to_gpu(multigrid)
    solution_gpu = zeros_like(x0_gpu)

    #CPU testing
    tic = timer.time()

    for _ in range(runs):
        (solution_cpu,
         r_norm_cpu,
         iterations_cpu) = amg_bicgstab_cpu(multigrid,b,x0,max_iterations=10,
                                            tol=tol,cycle='V')

    toc = timer.time()
    print('CPU time = ', toc - tic, ' s')
    print('relative residual = ' + str(r_norm_cpu))
    print('iterations = ' + str(iterations_cpu))


    #GPU testing
    tic = timer.time()
    for _ in range(runs):
        (solution_gpu,
         r_norm_gpu,
         iterations_gpu) = amg_bicgstab(multigrid_gpu,b_gpu,x0_gpu,max_iterations=10,
                                        tol=tol,cycle='V')

    toc = timer.time()
    print('GPU time = ', toc-tic,' s')
    print('relative residual = ' + str(r_norm_gpu))
    print('iterations = ' + str(iterations_gpu))


    plt.style.use('dark_background')

    plt.rcParams["font.family"] = "serif"
    plt.rcParams["font.serif"] = ["Times New Roman"]

    fig,(ax1,ax2) = plt.subplots(1,2,figsize=(12,5))

    ax1.imshow(solution_gpu.get().reshape((Nx,Ny)), cmap='turbo', interpolation='none',aspect="equal")
    ax1.set_xlabel('x (cells)',fontsize=17)
    ax1.set_ylabel('y (cells)',fontsize=17)
    ax1.tick_params(axis='both', labelsize=14)
    ax1.set_title('CUDA Solution',fontsize=20)


    ax2.imshow(solution_cpu.reshape((Nx, Ny)), cmap='turbo', interpolation='none', aspect="equal")
    ax2.set_xlabel('x (cells)',fontsize=17)
    ax2.set_ylabel('y (cells)',fontsize=17)
    ax2.tick_params(axis='both', labelsize=14)
    ax2.set_title('CPU Solution',fontsize=20)

    plt.show()