282 lines
11 KiB
Python
282 lines
11 KiB
Python
# 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.velo = np.zeros(self.FIELD_SIZE * self.DIM, 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 __del__(self):
|
|
self.context.pop() |