smyalygames/FiniteVolumeGPU
1
0
Fork 0
mirror of https://github.com/smyalygames/FiniteVolumeGPU.git synced 2026-01-14 15:48:43 +01:00
Code Issues Actions Packages Projects Releases Wiki Activity

Added class for managing SHMEMSimulators

Browse Source
This commit is contained in:
Martin Lilleeng Sætra
2021-05-05 15:24:34 +00:00
parent 2bc602622c
commit 2da4640840
3 changed files with 419 additions and 170 deletions
Download Patch File Download Diff File Expand all files Collapse all files

184
GPUSimulators/SHMEMSimulator.py
View File

@@ -21,139 +21,17 @@ along with this program. If not, see <http://www.gnu.org/licenses/>.
import logging
from GPUSimulators import Simulator, CudaContext
from GPUSimulators import Simulator
import numpy as np
import pycuda.driver as cuda
class SHMEMGrid(object):
"""
Class which represents an SHMEM grid of GPUs. Facilitates easy communication between
neighboring subdomains in the grid. Contains one CUDA context per subdomain.
"""
def __init__(self, ngpus=None, ndims=2):
self.logger = logging.getLogger(__name__)
cuda.init(flags=0)
self.logger.info("Initializing CUDA")
num_cuda_devices = cuda.Device.count()
if ngpus is None:
ngpus = num_cuda_devices
assert ngpus <= num_cuda_devices, "Trying to allocate more GPUs than are available in the system."
assert ndims == 2, "Unsupported number of dimensions. Must be two at the moment"
assert ngpus >= 2, "Must have at least two GPUs available to run multi-GPU simulations."
self.ngpus = ngpus
self.ndims = ndims
self.grid = SHMEMGrid.getGrid(self.ngpus, self.ndims)
self.logger.debug("Created {:}-dimensional SHMEM grid, using {:} GPUs".format(
self.ndims, self.ngpus))
self.cuda_contexts = []
for i in range(self.ngpus):
self.cuda_contexts.append(CudaContext.CudaContext(device=i, autotuning=False))
def getCoordinate(self, index):
i = (index % self.grid[0])
j = (index // self.grid[0])
return i, j
def getIndex(self, i, j):
return j*self.grid[0] + i
def getEast(self, index):
i, j = self.getCoordinate(index)
i = (i+1) % self.grid[0]
return self.getIndex(i, j)
def getWest(self, index):
i, j = self.getCoordinate(index)
i = (i+self.grid[0]-1) % self.grid[0]
return self.getIndex(i, j)
def getNorth(self, index):
i, j = self.getCoordinate(index)
j = (j+1) % self.grid[1]
return self.getIndex(i, j)
def getSouth(self, index):
i, j = self.getCoordinate(index)
j = (j+self.grid[1]-1) % self.grid[1]
return self.getIndex(i, j)
def getGrid(num_gpus, num_dims):
assert(isinstance(num_gpus, int))
assert(isinstance(num_dims, int))
# Adapted from https://stackoverflow.com/questions/28057307/factoring-a-number-into-roughly-equal-factors
# Original code by https://stackoverflow.com/users/3928385/ishamael
# Factorizes a number into n roughly equal factors
#Dictionary to remember already computed permutations
memo = {}
def dp(n, left): # returns tuple (cost, [factors])
"""
Recursively searches through all factorizations
"""
#Already tried: return existing result
if (n, left) in memo:
return memo[(n, left)]
#Spent all factors: return number itself
if left == 1:
return (n, [n])
#Find new factor
i = 2
best = n
bestTuple = [n]
while i * i < n:
#If factor found
if n % i == 0:
#Factorize remainder
rem = dp(n // i, left - 1)
#If new permutation better, save it
if rem[0] + i < best:
best = rem[0] + i
bestTuple = [i] + rem[1]
i += 1
#Store calculation
memo[(n, left)] = (best, bestTuple)
return memo[(n, left)]
grid = dp(num_gpus, num_dims)[1]
if (len(grid) < num_dims):
#Split problematic 4
if (4 in grid):
grid.remove(4)
grid.append(2)
grid.append(2)
#Pad with ones to guarantee num_dims
grid = grid + [1]*(num_dims - len(grid))
#Sort in descending order
grid = np.sort(grid)
grid = grid[::-1]
return grid
class SHMEMSimulator(Simulator.BaseSimulator):
"""
Class which handles communication between simulators on different GPUs
NOTE: This class is only intended to be used by SHMEMSimulatorGroup
"""
def __init__(self, sim, grid):
def __init__(self, index, sim, grid):
self.logger = logging.getLogger(__name__)
autotuner = sim.context.autotuner
@@ -168,14 +46,15 @@ class SHMEMSimulator(Simulator.BaseSimulator):
sim.block_size[0], sim.block_size[1])
sim.context.autotuner = autotuner
self.index = index
self.sim = sim
self.grid = grid
#Get neighbor subdomain ids
self.east = grid.getEast()
self.west = grid.getWest()
self.north = grid.getNorth()
self.south = grid.getSouth()
self.east = grid.getEast(self.index)
self.west = grid.getWest(self.index)
self.north = grid.getNorth(self.index)
self.south = grid.getSouth(self.index)
#Get coordinate of this subdomain
#and handle global boundary conditions
@@ -185,7 +64,7 @@ class SHMEMSimulator(Simulator.BaseSimulator):
'east': Simulator.BoundaryCondition.Type.Dirichlet,
'west': Simulator.BoundaryCondition.Type.Dirichlet
})
gi, gj = grid.getCoordinate()
gi, gj = grid.getCoordinate(self.index)
if (gi == 0 and boundary_conditions.west != Simulator.BoundaryCondition.Type.Periodic):
self.west = None
new_boundary_conditions.west = boundary_conditions.west;
@@ -237,7 +116,7 @@ class SHMEMSimulator(Simulator.BaseSimulator):
self.out_n = np.empty_like(self.in_n)
self.out_s = np.empty_like(self.in_s)
self.logger.debug("Simlator subdomain {:d} initialized on {:s}".format(self.grid.comm.rank, MPI.Get_processor_name()))
self.logger.debug("Simlator subdomain {:d} initialized on context {:s}".format(self.index, sim.context))
def substep(self, dt, step_number):
@@ -252,23 +131,22 @@ class SHMEMSimulator(Simulator.BaseSimulator):
def check(self):
return self.sim.check()
def computeDt(self):
local_dt = np.array([np.float32(self.sim.computeDt())]);
local_dt = np.array([np.float32(self.sim.computeDt())])
global_dt = np.empty(1, dtype=np.float32)
self.grid.comm.Allreduce(local_dt, global_dt, op=MPI.MIN)
self.logger.debug("Local dt: {:f}, global dt: {:f}".format(local_dt[0], global_dt[0]))
return global_dt[0]
def getExtent(self):
"""
Function which returns the extent of node with rank
rank in the grid
Function which returns the extent of the subdomain with index
index in the grid
"""
width = self.sim.nx*self.sim.dx
height = self.sim.ny*self.sim.dy
i, j = self.grid.getCoordinate()
i, j = self.grid.getCoordinate(self.index)
x0 = i * width
y0 = j * height
x1 = x0 + width
@@ -276,6 +154,19 @@ class SHMEMSimulator(Simulator.BaseSimulator):
return [x0, x1, y0, y1]
def exchange(self):
ns_download_before_exchange()
# GLOBAL SYNC
ns_do_exchange()
# GLOBAL SYNC
ns_upload_after_exchange()
ew_download_before_exchange()
# GLOBAL SYNC
ew_do_exchange()
# GLOBAL SYNC
ew_upload_after_exchange()
def ns_download_before_exchange(self):
####
# First transfer internal cells north-south
####
@@ -288,7 +179,8 @@ class SHMEMSimulator(Simulator.BaseSimulator):
for k in range(self.nvars):
self.sim.u0[k].download(self.sim.stream, cpu_data=self.out_s[k,:,:], asynch=True, extent=self.read_s)
self.sim.stream.synchronize()
def ns_do_exchange(self):
#Send/receive to north/south neighbours
comm_send = []
comm_recv = []
@@ -302,7 +194,8 @@ class SHMEMSimulator(Simulator.BaseSimulator):
#Wait for incoming transfers to complete
for comm in comm_recv:
comm.wait()
def ns_upload_after_exchange(self):
#Upload to the GPU
if self.north is not None:
for k in range(self.nvars):
@@ -314,9 +207,8 @@ class SHMEMSimulator(Simulator.BaseSimulator):
#Wait for sending to complete
for comm in comm_send:
comm.wait()
def ew_download_before_exchange(self):
####
# Then transfer east-west including ghost cells that have been filled in by north-south transfer above
####
@@ -329,7 +221,8 @@ class SHMEMSimulator(Simulator.BaseSimulator):
for k in range(self.nvars):
self.sim.u0[k].download(self.sim.stream, cpu_data=self.out_w[k,:,:], asynch=True, extent=self.read_w)
self.sim.stream.synchronize()
def ew_do_exchange(self):
#Send/receive to east/west neighbours
comm_send = []
comm_recv = []
@@ -344,7 +237,8 @@ class SHMEMSimulator(Simulator.BaseSimulator):
#Wait for incoming transfers to complete
for comm in comm_recv:
comm.wait()
def ew_upload_after_exchange(self):
#Upload to the GPU
if self.east is not None:
for k in range(self.nvars):