# CelerisLab/driver.py import pycuda.driver as cuda import numpy as np from typing import List, Tuple, Union from . import utils from . import preprocess as preproc from . import compiler FLUID = 0b00000001 SOLID = 0b00000010 GAS = 0b00000100 INTERFACE = 0b00001000 SENSOR = 0b00010000 class FlowField: def __init__( self, field_config: utils.FlowFieldConfig, cuda_config: utils.CudaConfig, device_id: Union[int, List[int]] = None, ): self.field_config = field_config self.cuda_config = cuda_config cuda.init() # Sanity checks if cuda_config.multi_gpu: if device_id is None or isinstance(device_id, int): raise ValueError("Multi-GPU support requires a list of device IDs.") # self.devices = [cuda.Device(id) for id in device_id] raise NotImplementedError("Multi-GPU support is not implemented yet.") else: if isinstance(device_id, list): if len(device_id) > 1: raise ValueError( "Single-GPU mode does not support multiple device IDs." ) device_id = device_id[0] elif device_id is None: device_id = 0 utils.check_cuda_device_availability(device_id) self.device = cuda.Device(device_id) self.context = self.device.make_context() utils.check_cuda_capability(field_config, cuda_config, device_id) # Config kernel compiler.config_kernal(cuda_config, field_config) compiler.config_object(int(0)) # compiler.config_sensor(int(0)) # Set constants if field_config.data_type == "FP32": self.DATA_TYPE = np.float32 else: raise ValueError(f"Unsupported data type {field_config.data_type}.") self.FIELD_SHAPE = tuple( size * unit for size, unit in zip( field_config.field_dim_in_U, cuda_config.unit_dimensions ) ) self.FIELD_SIZE = np.prod(self.FIELD_SHAPE) self.LATTICE = field_config.lattice self.DIM = field_config.dimensionality if field_config.lattice == 9 and field_config.dimensionality == 2: self.E = np.array( [0, 0, 1, 0, 0, 1, -1, 0, 0, -1, 1, 1, -1, 1, -1, -1, 1, -1], dtype=np.int32, ).reshape(9, 2) self.OPP = np.array([0, 3, 4, 1, 2, 7, 8, 5, 6], dtype=np.int32) self.WW = np.array( [4 / 9, 1 / 9, 1 / 9, 1 / 9, 1 / 9, 1 / 36, 1 / 36, 1 / 36, 1 / 36], dtype=self.DATA_TYPE, ) else: raise NotImplementedError( f"Unsupported lattice type {field_config.lattice} in {field_config.dimensionality} dimensions." ) # Compile kernel compiler.compile_kernel() self.ptx = cuda.module_from_file(compiler.kernel_path("kernel.ptx")) self.step = self.ptx.get_function("OneStep") initflow = self.ptx.get_function("InitTubeFlow") # Initialize memory self.ddf = np.zeros(self.FIELD_SIZE * self.LATTICE, dtype=self.DATA_TYPE) self.ddf_save = np.zeros(self.FIELD_SIZE * self.LATTICE, dtype=self.DATA_TYPE) self.flag = np.ones(self.FIELD_SIZE, dtype=np.uint8) self.indx = np.zeros(self.FIELD_SIZE, dtype=np.int32) self.delta_curve = np.zeros(0, dtype=self.DATA_TYPE) self.ddf_gpu = cuda.mem_alloc(self.ddf.nbytes) self.temp_gpu = cuda.mem_alloc(self.ddf.nbytes) self.flag_gpu = cuda.mem_alloc(self.flag.nbytes) self.indx_gpu = cuda.mem_alloc(self.indx.nbytes) self.objects = {} self.action = np.zeros(0, dtype=self.DATA_TYPE) self.obs = np.zeros(0, dtype=self.DATA_TYPE) initflow( self.flag_gpu, self.ddf_gpu, block=(self.cuda_config.threads_per_block, 1, 1), grid=( int(self.FIELD_SHAPE[0] / self.cuda_config.threads_per_block), int(self.FIELD_SHAPE[1]), int(self.FIELD_SHAPE[2]), ), ) cuda.memcpy_dtoh(self.flag, self.flag_gpu) cuda.memcpy_dtoh(self.ddf, self.ddf_gpu) def add_cylinder(self, center: Tuple[float, float, float], radius: float): x_c, y_c, z_c = center if ( x_c - radius <= 0 or x_c + radius >= self.FIELD_SHAPE[0] - 1 or y_c - radius <= 0 or y_c + radius >= self.FIELD_SHAPE[1] - 1 ): raise ValueError("Cylinder is out of bounds.") index = self.delta_curve.size if self.delta_curve.size > 0 else 0 if self.DATA_TYPE == np.float32: id_object = np.int32(len(self.objects)) else: raise ValueError(f"Unsupported data type {self.DATA_TYPE}.") for x in range(int(x_c - radius) - 1, int(x_c + radius) + 1): for y in range(int(y_c - radius) - 1, int(y_c + radius) + 1): if (x - x_c) ** 2 + (y - y_c) ** 2 < radius**2: k = x + y * self.FIELD_SHAPE[0] self.flag[k] = SOLID delta_temp = np.zeros(11, dtype=self.DATA_TYPE) delta_temp[0] = id_object.view(self.DATA_TYPE) for i in range(self.LATTICE): x_neb = x + self.E[i][0] y_neb = y + self.E[i][1] if (x_neb - x_c) ** 2 + (y_neb - y_c) ** 2 >= radius**2: self.flag[k] |= INTERFACE x_i, y_i = preproc.find_circle_intersection( x, y, x_neb, y_neb, x_c, y_c, radius ) d_neb = np.sqrt((x_i - x_neb) ** 2 + (y_i - y_neb) ** 2) delta_temp[i] = d_neb / np.sqrt( self.E[i][0] ** 2 + self.E[i][1] ** 2 ) if self.flag[k] & INTERFACE: delta_temp[9] = (y_c - y) / radius delta_temp[10] = (x - x_c) / radius self.delta_curve = np.concatenate( (self.delta_curve, delta_temp) ) self.indx[k] = index index += delta_temp.size self.objects[id_object] = { "type": "cylinder", "center": center, "radius": radius, } if hasattr(self, "delta_gpu"): self.delta_gpu.free() self.delta_gpu = cuda.mem_alloc(self.delta_curve.nbytes) self.action = np.zeros(len(self.objects), dtype=self.DATA_TYPE) if hasattr(self, "action_gpu"): self.action_gpu.free() self.action_gpu = cuda.mem_alloc(self.action.nbytes) self.obs = np.zeros(len(self.objects) * self.DIM, dtype=self.DATA_TYPE) if hasattr(self, "force_gpu"): self.obs_gpu.free() self.obs_gpu = cuda.mem_alloc(self.obs.nbytes) cuda.memcpy_htod(self.delta_gpu, self.delta_curve) cuda.memcpy_htod(self.flag_gpu, self.flag) cuda.memcpy_htod(self.indx_gpu, self.indx) compiler.config_object(len(self.objects)) compiler.compile_kernel() self.ptx = cuda.module_from_file(compiler.kernel_path("kernel.ptx")) self.step = self.ptx.get_function("OneStep") def add_sensor(self, center: Tuple[float, float, float], radius: float): x_c, y_c, z_c = center if ( x_c - radius <= 0 or x_c + radius >= self.FIELD_SHAPE[0] - 1 or y_c - radius <= 0 or y_c + radius >= self.FIELD_SHAPE[1] - 1 ): raise ValueError("Sensor is out of bounds.") id_object = len(self.objects) for x in range(int(x_c - radius) - 1, int(x_c + radius) + 1): for y in range(int(y_c - radius) - 1, int(y_c + radius) + 1): if (x - x_c) ** 2 + (y - y_c) ** 2 < radius**2: k = x + y * self.FIELD_SHAPE[0] self.flag[k] |= SENSOR self.indx[k] = id_object self.objects[id_object] = { "type": "sensor", "center": center, } self.action = np.zeros(len(self.objects), dtype=self.DATA_TYPE) if hasattr(self, "action_gpu"): self.action_gpu.free() self.action_gpu = cuda.mem_alloc(self.action.nbytes) self.obs = np.zeros(len(self.objects) * self.DIM, dtype=self.DATA_TYPE) if hasattr(self, "force_gpu"): self.obs_gpu.free() self.obs_gpu = cuda.mem_alloc(self.obs.nbytes) cuda.memcpy_htod(self.flag_gpu, self.flag) cuda.memcpy_htod(self.indx_gpu, self.indx) compiler.config_object(len(self.objects)) compiler.compile_kernel() self.ptx = cuda.module_from_file(compiler.kernel_path("kernel.ptx")) self.step = self.ptx.get_function("OneStep") def run(self, num_steps: int, action_target: np.ndarray): if ( action_target.size != len(self.objects) or action_target.dtype != self.DATA_TYPE ): raise ValueError("action data type or size does not match the objects.") weight = 0.1 stream = cuda.Stream() action_pinned = cuda.pagelocked_empty_like(self.action) action_pinned[:] = self.action obs_pinned = cuda.pagelocked_empty_like(self.obs) self.obs[:] = 0 for i in range(num_steps): action_pinned = (1 - weight) * action_pinned + weight * action_target cuda.memcpy_htod_async(self.action_gpu, action_pinned, stream) self.step( self.flag_gpu, self.ddf_gpu, self.temp_gpu, self.indx_gpu, self.delta_gpu, self.action_gpu, self.obs_gpu, block=(self.cuda_config.threads_per_block, 1, 1), grid=( int(self.FIELD_SHAPE[0] / self.cuda_config.threads_per_block), int(self.FIELD_SHAPE[1]), int(self.FIELD_SHAPE[2]), ), stream=stream, ) self.ddf_gpu, self.temp_gpu = self.temp_gpu, self.ddf_gpu cuda.memcpy_dtoh_async(obs_pinned, self.obs_gpu, stream) cuda.memset_d32_async(self.obs_gpu, 0, self.obs.size, stream) self.obs += obs_pinned stream.synchronize() self.obs = (self.obs / num_steps).astype(self.DATA_TYPE) def apply_ddf(self): cuda.memcpy_htod(self.ddf_gpu, self.ddf) def get_ddf(self): cuda.memcpy_dtoh(self.ddf, self.ddf_gpu) def save_ddf(self): self.ddf_save = self.ddf.copy() def restore_ddf(self): self.ddf = self.ddf_save.copy() def __del__(self): self.context.pop()