Major update: functional LBM and DRL
This commit is contained in:
Frank14f
parent
89f2678579
commit
c6f607dd8d
2
.gitignore
vendored
Normal file
2
.gitignore
vendored
Normal file
@ -0,0 +1,2 @@
|
||||
tensorboard/*
|
||||
models/*
|
||||
5
.vscode/settings.json
vendored
Normal file
5
.vscode/settings.json
vendored
Normal file
@ -0,0 +1,5 @@
|
||||
{
|
||||
"[cuda-cpp]": {
|
||||
},
|
||||
"C_Cpp.errorSquiggles": "disabled"
|
||||
}
|
||||
@ -1,86 +0,0 @@
|
||||
#include "setting.h"
|
||||
|
||||
__global__ void Collision(int *flag, LBtype *pres, LBtype *vell, LBtype *f0, LBtype *forc)
|
||||
{
|
||||
GLOBAL_INDEX()
|
||||
|
||||
if(flag[k]==0)
|
||||
{
|
||||
LBtype P,Ux,Uy;
|
||||
LBtype M[9];
|
||||
LBtype g[9];
|
||||
|
||||
LBtype Fx=forc[k*2], Fy=forc[k*2+1];
|
||||
|
||||
for(int kk=0;kk<9;kk++)
|
||||
g[kk]=f0[k*9+kk];
|
||||
|
||||
Ux=(g[1]+g[5]+g[8]-g[3]-g[6]-g[7]+0.5*Fx)/rho;
|
||||
Uy=(g[2]+g[5]+g[6]-g[4]-g[7]-g[8]+0.5*Fy)/rho;
|
||||
P =(g[0]+g[1]+g[2]+g[3]+g[4]+g[5]+g[6]+g[7]+g[8])/3.0;
|
||||
pressure[k]=P;
|
||||
|
||||
M[0]= g[0] +g[1] +g[2] +g[3] +g[4] +g[5] +g[6] +g[7] +g[8];
|
||||
M[1]=-4*g[0] -g[1] -g[2] -g[3] -g[4]+2*g[5]+2*g[6]+2*g[7]+2*g[8];
|
||||
M[2]= 4*g[0]-2*g[1]-2*g[2]-2*g[3]-2*g[4] +g[5] +g[6] +g[7] +g[8];
|
||||
M[3]= g[1] -g[3] +g[5] -g[6] -g[7] +g[8];
|
||||
M[4]= -2*g[1] +2*g[3] +g[5] -g[6] -g[7] +g[8];
|
||||
M[5]= g[2] -g[4] +g[5] +g[6] -g[7] -g[8];
|
||||
M[6]= -2*g[2] +2*g[4] +g[5] +g[6] -g[7] -g[8];
|
||||
M[7]= g[1] -g[2] +g[3] -g[4];
|
||||
M[8]= g[5] -g[6] +g[7] -g[8];
|
||||
|
||||
M[0]=1.00*( 3*P -M[0]);
|
||||
M[1]=1.20*(-6*P+3*rho*(Ux*Ux+Uy*Uy)-M[1])+(1-0.5*1.2)*6*(Ux*Fx+Uy*Fy);
|
||||
M[2]=1.20*( 3*P-3*rho*(Ux*Ux+Uy*Uy)-M[2])-(1-0.5*1.2)*6*(Ux*Fx+Uy*Fy);
|
||||
M[3]=1.00*( rho*Ux -M[3])+(1-0.5*0.0)*Fx;
|
||||
M[4]=1.20*(-rho*Ux -M[4])-(1-0.5*1.2)*Fx;
|
||||
M[5]=1.00*( rho*Uy -M[5])+(1-0.5*0.0)*Fy;
|
||||
M[6]=1.20*(-rho*Uy -M[6])-(1-0.5*1.2)*Fy;
|
||||
M[7]= nu*(rho*(Ux*Ux-Uy*Uy) -M[7])+(1-0.5*nu)*2*(Ux*Fx-Uy*Fy);
|
||||
M[8]= nu*(rho*Ux*Uy -M[8])+(1-0.5*nu)*(Ux*Fy+Uy*Fx);
|
||||
|
||||
f0[k*9] =g[0]+( M[0] -M[1] +M[2])/9.0;
|
||||
f0[k*9+1]=g[1]+(4*M[0] -M[1]-2*M[2]+6*M[3]-6*M[4] +9*M[7])/36.0;
|
||||
f0[k*9+2]=g[2]+(4*M[0] -M[1]-2*M[2] +6*M[5]-6*M[6]-9*M[7])/36.0;
|
||||
f0[k*9+3]=g[3]+(4*M[0] -M[1]-2*M[2]-6*M[3]+6*M[4] +9*M[7])/36.0;
|
||||
f0[k*9+4]=g[4]+(4*M[0] -M[1]-2*M[2] -6*M[5]+6*M[6]-9*M[7])/36.0;
|
||||
f0[k*9+5]=g[5]+(4*M[0]+2*M[1] +M[2]+6*M[3]+3*M[4]+6*M[5]+3*M[6] +9*M[8])/36.0;
|
||||
f0[k*9+6]=g[6]+(4*M[0]+2*M[1] +M[2]-6*M[3]-3*M[4]+6*M[5]+3*M[6] -9*M[8])/36.0;
|
||||
f0[k*9+7]=g[7]+(4*M[0]+2*M[1] +M[2]-6*M[3]-3*M[4]-6*M[5]-3*M[6] +9*M[8])/36.0;
|
||||
f0[k*9+8]=g[8]+(4*M[0]+2*M[1] +M[2]+6*M[3]+3*M[4]-6*M[5]-3*M[6] -9*M[8])/36.0;
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void Streaming(LBtype *f0, LBtype *f1)
|
||||
{
|
||||
GLOBAL_INDEX()
|
||||
int neighbor,nex,ney;
|
||||
int e[9][2]={{0,0},{1,0},{0,1},{-1,0},{0,-1},{1,1},{-1,1},{-1,-1},{1,-1}};
|
||||
|
||||
for(int kk=0;kk<9;kk++)
|
||||
{
|
||||
nex=(x+e[kk][0]+NX)%NX;
|
||||
ney=(y+e[kk][1]+NY)%NY;
|
||||
neighbor=ney*NX+nex;
|
||||
f1[neighbor*9+kk]=f0[k*9+kk];
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void BounceBack(int *flag, LBtype *f0)
|
||||
{
|
||||
GLOBAL_INDEX()
|
||||
int neighbor,nex,ney;
|
||||
int e[9][2]={{0,0},{1,0},{0,1},{-1,0},{0,-1},{1,1},{-1,1},{-1,-1},{1,-1}};
|
||||
int opp[9]={0,3,4,1,2,7,8,5,6};
|
||||
|
||||
if(flag[k]==1)
|
||||
for(int kk=1;kk<9;kk++)
|
||||
{
|
||||
nex=(x+e[kk][0]+NX)%NX;
|
||||
ney=(y+e[kk][1]+NY)%NY;
|
||||
neighbor=ney*NX+nex;
|
||||
if(flag[neighbor]==0)
|
||||
f0[neighbor*9+kk]=f0[k*9+opp[kk]];
|
||||
}
|
||||
}
|
||||
@ -1,9 +0,0 @@
|
||||
#define LBtype double
|
||||
#define Pi 3.141592653589793238
|
||||
const int N_thread=256;
|
||||
int devicenum=0;
|
||||
|
||||
#define GLOBAL_INDEX() \
|
||||
int x = threadIdx.x + blockDim.x * blockIdx.x; \
|
||||
int y = blockIdx.y; \
|
||||
int k = y * gridDim.x * blockDim.x + x;
|
||||
3
CelerisLab/__init__.py
Normal file
3
CelerisLab/__init__.py
Normal file
@ -0,0 +1,3 @@
|
||||
# CelerisLab/__init__.py
|
||||
|
||||
from .driver import FlowField
|
||||
BIN
CelerisLab/__pycache__/__init__.cpython-310.pyc
Normal file
BIN
CelerisLab/__pycache__/__init__.cpython-310.pyc
Normal file
Binary file not shown.
BIN
CelerisLab/__pycache__/compiler.cpython-310.pyc
Normal file
BIN
CelerisLab/__pycache__/compiler.cpython-310.pyc
Normal file
Binary file not shown.
BIN
CelerisLab/__pycache__/driver.cpython-310.pyc
Normal file
BIN
CelerisLab/__pycache__/driver.cpython-310.pyc
Normal file
Binary file not shown.
BIN
CelerisLab/__pycache__/preprocess.cpython-310.pyc
Normal file
BIN
CelerisLab/__pycache__/preprocess.cpython-310.pyc
Normal file
Binary file not shown.
BIN
CelerisLab/__pycache__/utils.cpython-310.pyc
Normal file
BIN
CelerisLab/__pycache__/utils.cpython-310.pyc
Normal file
Binary file not shown.
87
CelerisLab/compiler.py
Normal file
87
CelerisLab/compiler.py
Normal file
@ -0,0 +1,87 @@
|
||||
# CelerisLab/kernels/compiler.py
|
||||
|
||||
import subprocess
|
||||
import re
|
||||
import os
|
||||
|
||||
from .utils import FlowFieldConfig, CudaConfig
|
||||
|
||||
|
||||
def kernel_path(file_name: str) -> str:
|
||||
current_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
return os.path.join(current_dir, "kernels", file_name)
|
||||
|
||||
|
||||
def read_lines(file_path):
|
||||
with open(file_path, "r") as file:
|
||||
lines = file.readlines()
|
||||
return lines
|
||||
|
||||
|
||||
def write_lines(file_path, lines):
|
||||
with open(file_path, "w") as file:
|
||||
file.writelines(lines)
|
||||
|
||||
|
||||
def modify_macro(lines, macro_name, new_value):
|
||||
pattern = re.compile(rf"(#define\s+{macro_name}\s+)(\S+)")
|
||||
for i, line in enumerate(lines):
|
||||
if pattern.match(line):
|
||||
lines[i] = pattern.sub(rf"\g<1>{new_value}", line)
|
||||
break
|
||||
return lines
|
||||
|
||||
def modify_const(lines, const_name, new_type, new_value):
|
||||
pattern = re.compile(rf"(__constant__\s+)(\S+\s+{const_name}\s*=\s*)([^;]+)(;)")
|
||||
for i, line in enumerate(lines):
|
||||
if pattern.match(line):
|
||||
lines[i] = pattern.sub(rf"\g<1>{new_type} {const_name} = {new_value}\4", line)
|
||||
break
|
||||
return lines
|
||||
|
||||
def compile_kernel():
|
||||
subprocess.run(
|
||||
[
|
||||
"nvcc",
|
||||
"-ptx",
|
||||
kernel_path("kernel.cu"),
|
||||
"-o",
|
||||
kernel_path("kernel.ptx"),
|
||||
]
|
||||
)
|
||||
|
||||
def config_kernal(config_cuda: CudaConfig, config_field: FlowFieldConfig):
|
||||
lines = read_lines(kernel_path("macros.h"))
|
||||
lines = modify_macro(lines, "MULT_GPU", config_cuda.multi_gpu)
|
||||
lines = modify_macro(lines, "NT", config_cuda.threads_per_block)
|
||||
lines = modify_macro(lines, "X_1U", config_cuda.unit_dimensions[0])
|
||||
lines = modify_macro(lines, "Y_1U", config_cuda.unit_dimensions[1])
|
||||
lines = modify_macro(lines, "Z_1U", config_cuda.unit_dimensions[2])
|
||||
|
||||
if config_field.data_type == "FP32":
|
||||
lb_type = "float"
|
||||
else:
|
||||
raise ValueError(f"Unsupported data type {config_field.data_type}.")
|
||||
lines = modify_macro(lines, "LBtype", lb_type)
|
||||
lines = modify_macro(lines, "UX", config_field.field_dim_in_U[0])
|
||||
lines = modify_macro(lines, "UY", config_field.field_dim_in_U[1])
|
||||
lines = modify_macro(lines, "UZ", config_field.field_dim_in_U[2])
|
||||
lines = modify_macro(lines, "NX", config_field.field_dim_in_U[0] * config_cuda.unit_dimensions[0])
|
||||
lines = modify_macro(lines, "NY", config_field.field_dim_in_U[1] * config_cuda.unit_dimensions[1])
|
||||
lines = modify_macro(lines, "NZ", config_field.field_dim_in_U[2] * config_cuda.unit_dimensions[2])
|
||||
lines = modify_macro(lines, "DIM", config_field.dimensionality)
|
||||
lines = modify_macro(lines, "NQ", config_field.lattice)
|
||||
lines = modify_macro(lines, "VIS", config_field.viscosity)
|
||||
lines = modify_macro(lines, "U0", config_field.velocity)
|
||||
|
||||
write_lines(kernel_path("macros.h"), lines)
|
||||
|
||||
def config_object(n_obj: int):
|
||||
lines = read_lines(kernel_path("macros.h"))
|
||||
lines = modify_macro(lines, "N_OBJS", n_obj)
|
||||
write_lines(kernel_path("macros.h"), lines)
|
||||
|
||||
def config_sensor(n_sen: int):
|
||||
lines = read_lines(kernel_path("macros.h"))
|
||||
lines = modify_macro(lines, "N_SENS", n_sen)
|
||||
write_lines(kernel_path("macros.h"), lines)
|
||||
282
CelerisLab/driver.py
Normal file
282
CelerisLab/driver.py
Normal file
@ -0,0 +1,282 @@
|
||||
# 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()
|
||||
101
CelerisLab/kernels/D2Q9.cu
Normal file
101
CelerisLab/kernels/D2Q9.cu
Normal file
@ -0,0 +1,101 @@
|
||||
#include "macros.h"
|
||||
#include "const.h"
|
||||
|
||||
__device__ void Index_lattice(int &x, int &y, int &k) {
|
||||
// Only for D2
|
||||
x = threadIdx.x + NT * blockIdx.x;
|
||||
y = blockIdx.y;
|
||||
k = y * NX + x;
|
||||
}
|
||||
|
||||
__device__ void CollisionKernel(LBtype* g, LBtype* m) {
|
||||
// Only for D2Q9
|
||||
LBtype p, u, v;
|
||||
LBtype niu = 1.0 / (0.5 + 3 * VIS);
|
||||
|
||||
u = (g[1]+g[5]+g[8]-g[3]-g[6]-g[7])/RHO;
|
||||
v = (g[2]+g[5]+g[6]-g[4]-g[7]-g[8])/RHO;
|
||||
p = (g[0]+g[1]+g[2]+g[3]+g[4]+g[5]+g[6]+g[7]+g[8])/3.0;
|
||||
|
||||
m[0]= g[0] +g[1] +g[2] +g[3] +g[4] +g[5] +g[6] +g[7] +g[8];
|
||||
m[1]=-4*g[0] -g[1] -g[2] -g[3] -g[4]+2*g[5]+2*g[6]+2*g[7]+2*g[8];
|
||||
m[2]= 4*g[0]-2*g[1]-2*g[2]-2*g[3]-2*g[4] +g[5] +g[6] +g[7] +g[8];
|
||||
m[3]= g[1] -g[3] +g[5] -g[6] -g[7] +g[8];
|
||||
m[4]= -2*g[1] +2*g[3] +g[5] -g[6] -g[7] +g[8];
|
||||
m[5]= g[2] -g[4] +g[5] +g[6] -g[7] -g[8];
|
||||
m[6]= -2*g[2] +2*g[4] +g[5] +g[6] -g[7] -g[8];
|
||||
m[7]= g[1] -g[2] +g[3] -g[4];
|
||||
m[8]= g[5] -g[6] +g[7] -g[8];
|
||||
|
||||
m[0]=1.00*( 3*p -m[0]);
|
||||
m[1]=1.20*(-6*p +3*RHO*(u*u+v*v)-m[1]);
|
||||
m[2]=1.20*( 3*p -3*RHO*(u*u+v*v)-m[2]);
|
||||
m[3]=1.00*( RHO*u -m[3]);
|
||||
m[4]=1.20*(-RHO*u -m[4]);
|
||||
m[5]=1.00*( RHO*v -m[5]);
|
||||
m[6]=1.20*(-RHO*v -m[6]);
|
||||
m[7]= niu*( RHO*(u*u-v*v) -m[7]);
|
||||
m[8]= niu*( RHO*u*v -m[8]);
|
||||
|
||||
g[0]=g[0]+( m[0] -m[1] +m[2] )/ 9.0;
|
||||
g[1]=g[1]+(4*m[0] -m[1]-2*m[2]+6*m[3]-6*m[4] +9*m[7])/36.0;
|
||||
g[2]=g[2]+(4*m[0] -m[1]-2*m[2] +6*m[5]-6*m[6]-9*m[7])/36.0;
|
||||
g[3]=g[3]+(4*m[0] -m[1]-2*m[2]-6*m[3]+6*m[4] +9*m[7])/36.0;
|
||||
g[4]=g[4]+(4*m[0] -m[1]-2*m[2] -6*m[5]+6*m[6]-9*m[7])/36.0;
|
||||
g[5]=g[5]+(4*m[0]+2*m[1] +m[2]+6*m[3]+3*m[4]+6*m[5]+3*m[6]+9*m[8])/36.0;
|
||||
g[6]=g[6]+(4*m[0]+2*m[1] +m[2]-6*m[3]-3*m[4]+6*m[5]+3*m[6]-9*m[8])/36.0;
|
||||
g[7]=g[7]+(4*m[0]+2*m[1] +m[2]-6*m[3]-3*m[4]-6*m[5]-3*m[6]+9*m[8])/36.0;
|
||||
g[8]=g[8]+(4*m[0]+2*m[1] +m[2]+6*m[3]+3*m[4]-6*m[5]-3*m[6]-9*m[8])/36.0;
|
||||
}
|
||||
|
||||
__device__ void ParabolicInlet(LBtype* f, LBtype* f_neb, LBtype y) {
|
||||
LBtype p, u, v, yy;
|
||||
LBtype feq1, feq5, feq8, feqn1, feqn5, feqn8;
|
||||
|
||||
p=(f_neb[0]+f_neb[1]+f_neb[2]+f_neb[3]+f_neb[4]+f_neb[5]+f_neb[6]+f_neb[7]+f_neb[8])/3.0;
|
||||
yy=(y-0.5*(NY-1))/(NY-2.0);
|
||||
u=U0*1.5*(1-4*yy*yy);
|
||||
v=0.0;
|
||||
|
||||
feq1=(2*p+RHO*(2*u*u+2*u -v*v) )/ 6.0;
|
||||
feq5=( p+RHO*( u*u+3*u*v+u+v*v+v))/12.0;
|
||||
feq8=( p+RHO*( u*u-3*u*v+u+v*v-v))/12.0;
|
||||
|
||||
u=(f_neb[1]+f_neb[5]+f_neb[8]-f_neb[3]-f_neb[6]-f_neb[7])/RHO;
|
||||
v=(f_neb[2]+f_neb[5]+f_neb[6]-f_neb[4]-f_neb[7]-f_neb[8])/RHO;
|
||||
|
||||
feqn1=(2*p+RHO*(2*u*u+2*u -v*v) )/ 6.0;
|
||||
feqn5=( p+RHO*( u*u+3*u*v+u+v*v+v))/12.0;
|
||||
feqn8=( p+RHO*( u*u-3*u*v+u+v*v-v))/12.0;
|
||||
|
||||
f[1]=f_neb[1]-feqn1+feq1;
|
||||
f[5]=f_neb[5]-feqn5+feq5;
|
||||
f[8]=f_neb[8]-feqn8+feq8;
|
||||
}
|
||||
|
||||
__device__ void PressureOutlet(LBtype* f, LBtype* f_neb, LBtype y) {
|
||||
// Edit to Parabolic Outlet temporarily
|
||||
LBtype p, u, v, yy;
|
||||
LBtype feq3, feq6, feq7, feqn3, feqn6, feqn7;
|
||||
|
||||
p=0.0;
|
||||
|
||||
yy=(y-0.5*(NY-1))/(NY-2.0);
|
||||
u=U0*1.5*(1-4*yy*yy);
|
||||
v=0.0;
|
||||
|
||||
feq3=(2*p-RHO*(-2*u*u+2*u +v*v) )/ 6.0;
|
||||
feq6=( p+RHO*( u*u-3*u*v-u+v*v+v))/12.0;
|
||||
feq7=( p+RHO*( u*u+3*u*v-u+v*v-v))/12.0;
|
||||
|
||||
u=(f_neb[1]+f_neb[5]+f_neb[8]-f_neb[3]-f_neb[6]-f_neb[7])/RHO;
|
||||
v=(f_neb[2]+f_neb[5]+f_neb[6]-f_neb[4]-f_neb[7]-f_neb[8])/RHO;
|
||||
// p=(f_neb[0]+f_neb[1]+f_neb[2]+f_neb[3]+f_neb[4]+f_neb[5]+f_neb[6]+f_neb[7]+f_neb[8])/3.0;
|
||||
feqn3=(2*p-RHO*(-2*u*u+2*u +v*v) )/ 6.0;
|
||||
feqn6=( p+RHO*( u*u-3*u*v-u+v*v+v))/12.0;
|
||||
feqn7=( p+RHO*( u*u+3*u*v-u+v*v-v))/12.0;
|
||||
|
||||
f[3]=f_neb[3]-feqn3+feq3;
|
||||
f[6]=f_neb[6]-feqn6+feq6;
|
||||
f[7]=f_neb[7]-feqn7+feq7;
|
||||
}
|
||||
0
CelerisLab/kernels/IO.cu
Normal file
0
CelerisLab/kernels/IO.cu
Normal file
BIN
CelerisLab/kernels/__pycache__/compiler.cpython-310.pyc
Normal file
BIN
CelerisLab/kernels/__pycache__/compiler.cpython-310.pyc
Normal file
Binary file not shown.
10
CelerisLab/kernels/const.h
Normal file
10
CelerisLab/kernels/const.h
Normal file
@ -0,0 +1,10 @@
|
||||
// CelerisLab/kernels/const.h
|
||||
|
||||
#ifndef CONST_H
|
||||
#define CONST_H
|
||||
|
||||
__constant__ int e[9][2] = {{0, 0}, {1, 0}, {0, 1}, {-1, 0}, {0, -1}, {1, 1}, {-1, 1}, {-1, -1}, {1, -1}};
|
||||
__constant__ int opp[9] = {0, 3, 4, 1, 2, 7, 8, 5, 6};
|
||||
__constant__ float w[9] = {4/9., 1/9., 1/9., 1/9., 1/9., 1/36., 1/36., 1/36., 1/36.};
|
||||
|
||||
#endif
|
||||
190
CelerisLab/kernels/kernel.cu
Normal file
190
CelerisLab/kernels/kernel.cu
Normal file
@ -0,0 +1,190 @@
|
||||
// CelerisLab/kernels/kernel.cu
|
||||
|
||||
#include <stdio.h>
|
||||
#include <stdint.h>
|
||||
#include <cuda.h>
|
||||
|
||||
#include "macros.h"
|
||||
#include "const.h"
|
||||
#include "D2Q9.cu"
|
||||
|
||||
extern "C"
|
||||
{
|
||||
__global__ void OneStep(uint8_t *flag, LBtype *f, LBtype *f_temp, int32_t *indx, LBtype *delta, LBtype *action, LBtype *obs)
|
||||
{
|
||||
__shared__ LBtype f_share[NT * NQ];
|
||||
__shared__ LBtype obs_share[(N_OBJS * DIM > 0) ? N_OBJS * DIM : 1];
|
||||
|
||||
int x, y, k;
|
||||
LBtype g[NQ], m[NQ];
|
||||
Index_lattice(x, y, k); // Only for D2
|
||||
int totalCells = NX * NY;
|
||||
int id = indx[k];
|
||||
|
||||
for (int i = 0; i < NQ; i++)
|
||||
{
|
||||
f_share[threadIdx.x + i * NT] = f[k + i * totalCells];
|
||||
}
|
||||
for (int i = threadIdx.x; i < N_OBJS * DIM; i+=NT)
|
||||
{
|
||||
obs_share[i] = 0;
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
for (int i = 0; i < NQ; i++)
|
||||
{
|
||||
g[i] = f_share[threadIdx.x + i * NT];
|
||||
}
|
||||
|
||||
if (flag[k] & FLUID)
|
||||
{
|
||||
CollisionKernel(g, m);
|
||||
|
||||
for (int i = 0; i < NQ; i++)
|
||||
{
|
||||
f_share[threadIdx.x + i * NT] = g[i];
|
||||
}
|
||||
}
|
||||
else if (flag[k] & SOLID)
|
||||
{
|
||||
if (x == 0)
|
||||
{
|
||||
for (int i = 0; i < NQ; i++)
|
||||
{
|
||||
m[i] = f_share[threadIdx.x + i * NT + 1];
|
||||
}
|
||||
ParabolicInlet(g, m, y);
|
||||
}
|
||||
else if (x == NX - 1)
|
||||
{
|
||||
for (int i = 0; i < NQ; i++)
|
||||
{
|
||||
m[i] = f_share[threadIdx.x + i * NT - 1];
|
||||
}
|
||||
PressureOutlet(g, m, y);
|
||||
}
|
||||
|
||||
for (int i = 0; i < NQ; i++)
|
||||
{
|
||||
f_share[threadIdx.x + i * NT] = g[i];
|
||||
}
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
for (int i = 0; i < NQ; i++)
|
||||
{
|
||||
int x_neb = x + e[i][0];
|
||||
int y_neb = y + e[i][1];
|
||||
|
||||
if (y != 0 && y != NY - 1)
|
||||
{
|
||||
if ((y == 1 && y_neb == 0) || (y == NY - 2 && y_neb == NY - 1))
|
||||
{
|
||||
f_temp[k + opp[i] * totalCells] = f_share[threadIdx.x + i * NT];
|
||||
}
|
||||
else
|
||||
{
|
||||
int k_neb = ((y_neb * NX + x_neb) + totalCells) % totalCells;
|
||||
f_temp[k_neb + i * totalCells] = f_share[threadIdx.x + i * NT];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
if (flag[k] & SOLID && flag[k] & INTERFACE)
|
||||
{
|
||||
LBtype Uw, Vw;
|
||||
int id_obj = *reinterpret_cast<int*>(&delta[id]);
|
||||
Uw = action[id_obj] * delta[id + 9];
|
||||
Vw = action[id_obj] * delta[id + 10];
|
||||
|
||||
int x_neb, y_neb, k_neb;
|
||||
for (int i = 1; i < 9; i++)
|
||||
{
|
||||
x_neb = x + e[i][0];
|
||||
y_neb = y + e[i][1];
|
||||
k_neb = x_neb + y_neb * NX;
|
||||
if (flag[k_neb] & FLUID)
|
||||
{
|
||||
LBtype q = delta[id + i];
|
||||
int k_neb2 = (y + 2 * e[i][1]) * NX + (x + 2 * e[i][0]);
|
||||
LBtype temp = 6 * w[i] * (e[i][0] * Uw + e[i][1] * Vw);
|
||||
f_temp[k_neb + i * totalCells] = (q * f_temp[k + opp[i] * totalCells] \
|
||||
+ (1 - q) * f_temp[k_neb + opp[i] * totalCells] \
|
||||
+ q * f_temp[k_neb2 + i * totalCells] + temp) / (1 + q);
|
||||
f_temp[k + i * totalCells] = temp * Uw;
|
||||
k_neb2 = (y - e[i][1]) * NX + (x - e[i][0]);
|
||||
f_temp[k_neb2 + i * totalCells] = temp * Vw;
|
||||
|
||||
temp = f_temp[k_neb + i * totalCells] + f_temp[k + opp[i] * totalCells];
|
||||
k_neb2 = (y - e[i][1]) * NX + (x - e[i][0]);
|
||||
atomicAdd(&obs_share[DIM * id_obj], -temp * e[i][0] + f_temp[k + i * totalCells]);
|
||||
atomicAdd(&obs_share[DIM * id_obj + 1], -temp * e[i][1] + f_temp[k_neb2 + i * totalCells]);
|
||||
}
|
||||
}
|
||||
}
|
||||
if (flag[k] & SENSOR)
|
||||
{
|
||||
LBtype u, v;
|
||||
u = (g[1]+g[5]+g[8]-g[3]-g[6]-g[7])/RHO;
|
||||
v = (g[2]+g[5]+g[6]-g[4]-g[7]-g[8])/RHO;
|
||||
atomicAdd(&obs_share[DIM * id], u);
|
||||
atomicAdd(&obs_share[DIM * id + 1], v);
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
for (int i = threadIdx.x; i < N_OBJS * DIM; i+=NT)
|
||||
{
|
||||
atomicAdd(&obs[i], obs_share[i]);
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void InitTubeFlow(uint8_t *flag, LBtype *f)
|
||||
{
|
||||
__shared__ LBtype f_share[NT * NQ];
|
||||
__shared__ uint8_t flag_share[NT];
|
||||
int x, y, k;
|
||||
LBtype u;
|
||||
Index_lattice(x, y, k);
|
||||
int totalCells = NX * NY;
|
||||
|
||||
flag_share[threadIdx.x] = flag[k];
|
||||
for (int i = 0; i < NQ; i++)
|
||||
{
|
||||
f_share[threadIdx.x + i * NT] = f[k + i * totalCells];
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
u = U0 * 1.5 * (1 - 4 * (y - 0.5 * (NY - 1)) * (y - 0.5 * (NY - 1)) / ((NY - 2) * (NY - 2)));
|
||||
if (y == 0 || y == NY - 1 || x == 0 || x == NX - 1)
|
||||
{
|
||||
flag_share[threadIdx.x] = SOLID;
|
||||
for (int i = 0; i < NQ; i++)
|
||||
{
|
||||
f_share[threadIdx.x + i * NT] = 0;
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
flag_share[threadIdx.x] = FLUID;
|
||||
for (int i = 0; i < NQ; i++)
|
||||
{
|
||||
f_share[threadIdx.x + i * NT] = w[i] * RHO * (3 * e[i][0] * u + \
|
||||
4.5 * e[i][0] * e[i][0] * u * u - 1.5 * u * u);
|
||||
}
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
flag[k] = flag_share[threadIdx.x];
|
||||
for (int i = 0; i < NQ; i++)
|
||||
{
|
||||
f[k + i * totalCells] = f_share[threadIdx.x + i * NT];
|
||||
}
|
||||
}
|
||||
}
|
||||
1410
CelerisLab/kernels/kernel.ptx
Normal file
1410
CelerisLab/kernels/kernel.ptx
Normal file
File diff suppressed because it is too large
Load Diff
34
CelerisLab/kernels/macros.h
Normal file
34
CelerisLab/kernels/macros.h
Normal file
@ -0,0 +1,34 @@
|
||||
// CelerisLab/kernels/macros.h
|
||||
|
||||
// cuda parameters
|
||||
#define MULT_GPU False
|
||||
#define NT 128
|
||||
#define X_1U 128
|
||||
#define Y_1U 32
|
||||
#define Z_1U 1
|
||||
|
||||
// flow parameters
|
||||
#define LBtype float
|
||||
#define UX 10
|
||||
#define UY 16
|
||||
#define UZ 1
|
||||
#define NX 1280
|
||||
#define NY 512
|
||||
#define NZ 1
|
||||
#define DIM 2
|
||||
#define NQ 9
|
||||
#define VIS 0.004
|
||||
#define RHO 1.0
|
||||
#define U0 0.01
|
||||
|
||||
// constants
|
||||
#define PI 3.141592653589793238
|
||||
#define FLUID 0b00000001
|
||||
#define SOLID 0b00000010
|
||||
#define GAS 0b00000100
|
||||
#define INTERFACE 0b00001000
|
||||
#define SENSOR 0b00010000
|
||||
|
||||
// variables
|
||||
#define N_OBJS 7
|
||||
// #define N_SENS 2
|
||||
2
CelerisLab/kernels/preproc.cu
Normal file
2
CelerisLab/kernels/preproc.cu
Normal file
@ -0,0 +1,2 @@
|
||||
#include "macros.h"
|
||||
#include "const.h"
|
||||
40
CelerisLab/preprocess.py
Normal file
40
CelerisLab/preprocess.py
Normal file
@ -0,0 +1,40 @@
|
||||
# CelerisLab/preprocess.py
|
||||
|
||||
import math
|
||||
import numpy as np
|
||||
from typing import Tuple
|
||||
|
||||
FLUID = 0b00000001
|
||||
SOLID = 0b00000010
|
||||
GAS = 0b00000100
|
||||
INTERFACE = 0b00001000
|
||||
SENSOR = 0b00010000
|
||||
|
||||
|
||||
def find_circle_intersection(x, y, x_neb, y_neb, xc, yc, r0):
|
||||
dx, dy = x_neb - x, y_neb - y
|
||||
a = dx ** 2 + dy ** 2
|
||||
b = 2 * (dx * (x - xc) + dy * (y - yc))
|
||||
c = (x - xc) ** 2 + (y - yc) ** 2 - r0 ** 2
|
||||
det = b ** 2 - 4 * a * c
|
||||
|
||||
if det < 0:
|
||||
return None
|
||||
|
||||
t1 = (-b + math.sqrt(det)) / (2 * a)
|
||||
t2 = (-b - math.sqrt(det)) / (2 * a)
|
||||
|
||||
if 0 <= t1 <= 1:
|
||||
return x + t1 * dx, y + t1 * dy
|
||||
elif 0 <= t2 <= 1:
|
||||
return x + t2 * dx, y + t2 * dy
|
||||
else:
|
||||
return None
|
||||
|
||||
def find_sensor_area(radius):
|
||||
area = 0
|
||||
for i in range(np.floor(-radius), np.ceil(radius)):
|
||||
for j in range(np.floor(-radius), np.ceil(radius)):
|
||||
if i ** 2 + j ** 2 <= radius ** 2:
|
||||
area += 1
|
||||
return area
|
||||
256
CelerisLab/utils.py
Normal file
256
CelerisLab/utils.py
Normal file
@ -0,0 +1,256 @@
|
||||
# CelerisLab/utils.py
|
||||
|
||||
import pycuda.driver as cuda
|
||||
import subprocess
|
||||
import json
|
||||
|
||||
from typing import NamedTuple, Optional, List, Tuple, Union
|
||||
|
||||
|
||||
class CudaDeviceInfo(NamedTuple):
|
||||
name: str
|
||||
compute_capability: str
|
||||
multiprocessors: int
|
||||
total_global_memory: int
|
||||
max_shared_memory_per_block: int
|
||||
max_threads_per_block: int
|
||||
max_blocks_per_multiprocessor: int
|
||||
device_interconnect: Optional[str] = None
|
||||
|
||||
|
||||
class FlowFieldConfig(NamedTuple):
|
||||
data_type: str
|
||||
dimensionality: int
|
||||
lattice: int
|
||||
field_dim_in_U: Tuple[int, int, int]
|
||||
viscosity: float
|
||||
velocity: float
|
||||
boundary_conditions: Tuple[str, str, str, str, str, str]
|
||||
|
||||
|
||||
class CudaConfig(NamedTuple):
|
||||
multi_gpu: bool
|
||||
gpu_connection: str
|
||||
required_cuda_capability: str
|
||||
threads_per_block: int
|
||||
unit_dimensions: Tuple[int, int, int]
|
||||
|
||||
|
||||
def check_cuda_device_availability(device_id=0):
|
||||
if cuda.Device.count() == 0:
|
||||
raise RuntimeError("No CUDA device is available.")
|
||||
|
||||
if device_id < 0 or device_id >= cuda.Device.count():
|
||||
raise ValueError(
|
||||
f"Invalid device_id {device_id}. Must be between 0 and {cuda.Device.count() - 1}."
|
||||
)
|
||||
|
||||
try:
|
||||
subprocess.check_output(["nvidia-smi", "--version"])
|
||||
except subprocess.CalledProcessError:
|
||||
raise RuntimeError("nvidia-smi is not available or not installed correctly.")
|
||||
|
||||
|
||||
def query_cuda_device_info(device_id=0) -> CudaDeviceInfo:
|
||||
check_cuda_device_availability(device_id)
|
||||
|
||||
try:
|
||||
output = subprocess.check_output(
|
||||
["nvidia-smi", "-q", "-d", "TOPOLOGY", "-i", str(device_id)], text=True
|
||||
)
|
||||
if "NVLink" in output:
|
||||
device_interconnect = "NVLink"
|
||||
elif "PCIe" in output:
|
||||
device_interconnect = "PCIe"
|
||||
else:
|
||||
device_interconnect = "Unknown"
|
||||
except Exception as e:
|
||||
device_interconnect = None
|
||||
|
||||
device = cuda.Device(device_id)
|
||||
|
||||
return CudaDeviceInfo(
|
||||
name=device.name(),
|
||||
compute_capability=f"{device.compute_capability()[0]}.{device.compute_capability()[1]}",
|
||||
multiprocessors=device.get_attribute(
|
||||
cuda.device_attribute.MULTIPROCESSOR_COUNT
|
||||
),
|
||||
total_global_memory=device.total_memory(),
|
||||
max_shared_memory_per_block=device.get_attribute(
|
||||
cuda.device_attribute.MAX_SHARED_MEMORY_PER_BLOCK
|
||||
),
|
||||
max_threads_per_block=device.get_attribute(
|
||||
cuda.device_attribute.MAX_THREADS_PER_BLOCK
|
||||
),
|
||||
max_blocks_per_multiprocessor=device.get_attribute(
|
||||
cuda.device_attribute.MAX_BLOCKS_PER_MULTIPROCESSOR
|
||||
),
|
||||
device_interconnect=device_interconnect,
|
||||
)
|
||||
|
||||
|
||||
def load_flow_field_config(config_path: str) -> FlowFieldConfig:
|
||||
try:
|
||||
with open(config_path, "r") as file:
|
||||
config = json.load(file)
|
||||
|
||||
required_keys = [
|
||||
"data_type",
|
||||
"dimensionality",
|
||||
"lattice",
|
||||
"field_dim_in_U",
|
||||
"viscosity",
|
||||
"boundary_conditions",
|
||||
]
|
||||
if not all(key in config for key in required_keys):
|
||||
raise ValueError("Missing required configuration items.")
|
||||
|
||||
if config["data_type"] not in ["FP32", "FP64"]:
|
||||
raise ValueError("Data type must be either FP32 or FP64.")
|
||||
|
||||
if config["dimensionality"] not in [2, 3]:
|
||||
raise ValueError("Dimensionality must be either 2 or 3.")
|
||||
|
||||
if config["dimensionality"] == 2 and config["field_dim_in_U"][2] != 1:
|
||||
raise ValueError(
|
||||
"Field dimensions must be 1 in the third dimension for 2D simulations."
|
||||
)
|
||||
|
||||
if config["lattice"] not in [9]:
|
||||
raise ValueError("Lattice must be either 9 or 19.")
|
||||
|
||||
boundary_conditions = tuple(
|
||||
condition
|
||||
for key in ["x", "y", "z"]
|
||||
for condition in config["boundary_conditions"].get(key, [])
|
||||
)
|
||||
if len(boundary_conditions) != 6:
|
||||
raise ValueError("Boundary conditions must contain exactly six elements.")
|
||||
|
||||
return FlowFieldConfig(
|
||||
data_type=config["data_type"],
|
||||
dimensionality=config["dimensionality"],
|
||||
lattice=config["lattice"],
|
||||
field_dim_in_U=tuple(config["field_dim_in_U"]),
|
||||
viscosity=config["viscosity"],
|
||||
velocity=config["velocity"],
|
||||
boundary_conditions=boundary_conditions,
|
||||
)
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to load or parse the flow field configuration: {e}")
|
||||
|
||||
|
||||
def load_cuda_config(config_path: str) -> CudaConfig:
|
||||
try:
|
||||
with open(config_path, "r") as file:
|
||||
config = json.load(file)
|
||||
|
||||
required_keys = [
|
||||
"multi_gpu",
|
||||
"gpu_connection",
|
||||
"required_cuda_capability",
|
||||
"threads_per_block",
|
||||
"X_1U",
|
||||
"Y_1U",
|
||||
"Z_1U",
|
||||
]
|
||||
|
||||
if not all(key in config for key in required_keys):
|
||||
raise ValueError("Missing required configuration items.")
|
||||
|
||||
return CudaConfig(
|
||||
multi_gpu=config["multi_gpu"],
|
||||
gpu_connection=config["gpu_connection"],
|
||||
required_cuda_capability=config["required_cuda_capability"],
|
||||
threads_per_block=config["threads_per_block"],
|
||||
unit_dimensions=(config["X_1U"], config["Y_1U"], config["Z_1U"]),
|
||||
)
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to load or parse the CUDA configuration: {e}")
|
||||
|
||||
|
||||
def check_cuda_capability(
|
||||
field_config: FlowFieldConfig,
|
||||
cuda_config: CudaConfig,
|
||||
device_id: Union[int, List[int]] = None,
|
||||
):
|
||||
SAFE_FACTOR = 0.8
|
||||
|
||||
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.")
|
||||
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
|
||||
device_info = query_cuda_device_info(device_id)
|
||||
|
||||
if device_info.compute_capability != cuda_config.required_cuda_capability:
|
||||
raise ValueError(
|
||||
f"Device {device_info.name} has compute capability {device_info.compute_capability}, but {cuda_config.required_cuda_capability} is required."
|
||||
)
|
||||
|
||||
field_size = sum(
|
||||
size * unit
|
||||
for size, unit in zip(
|
||||
field_config.field_dim_in_U, cuda_config.unit_dimensions
|
||||
)
|
||||
)
|
||||
if (
|
||||
device_info.total_global_memory * SAFE_FACTOR
|
||||
< calc_field_memory_consumption(
|
||||
field_size,
|
||||
field_config.dimensionality,
|
||||
field_config.lattice,
|
||||
field_config.data_type,
|
||||
)
|
||||
):
|
||||
raise ValueError(
|
||||
f"Device {device_info.name} does not have enough memory to store the flow field."
|
||||
)
|
||||
|
||||
if (
|
||||
device_info.max_threads_per_block * SAFE_FACTOR
|
||||
< cuda_config.threads_per_block
|
||||
):
|
||||
raise ValueError(
|
||||
f"Device {device_info.name} does not have enough threads per block to run the simulation."
|
||||
)
|
||||
|
||||
block_size = cuda_config.threads_per_block
|
||||
if (
|
||||
device_info.max_shared_memory_per_block * SAFE_FACTOR
|
||||
< 2
|
||||
* calc_field_memory_consumption(
|
||||
block_size,
|
||||
field_config.dimensionality,
|
||||
field_config.lattice,
|
||||
field_config.data_type,
|
||||
)
|
||||
):
|
||||
raise ValueError(
|
||||
f"Device {device_info.name} does not have enough shared memory per block to run the simulation."
|
||||
)
|
||||
|
||||
|
||||
def calc_field_memory_consumption(
|
||||
field_size: int, dimensionality: int, directions: int, data_type: str
|
||||
) -> int:
|
||||
if data_type == "FP32":
|
||||
data_size = 4
|
||||
elif data_type == "FP64":
|
||||
data_size = 8
|
||||
else:
|
||||
raise ValueError(f"Unsupported data type {data_type}.")
|
||||
|
||||
return (
|
||||
field_size * directions * data_size * 2
|
||||
+ field_size * dimensionality * data_size
|
||||
+ field_size
|
||||
)
|
||||
Binary file not shown.
@ -1,74 +0,0 @@
|
||||
import gymnasium as gym
|
||||
import numpy as np
|
||||
from gymnasium import spaces
|
||||
import ctypes
|
||||
from collections import deque
|
||||
|
||||
lbm = ctypes.cdll.LoadLibrary('./lbm_sens.so')
|
||||
|
||||
S_DIM, A_DIM = 6, 3
|
||||
action_amp = 5
|
||||
action_weight = 0.5
|
||||
sample_interval = 200
|
||||
max_steps = 320
|
||||
|
||||
class CustomEnv(gym.Env):
|
||||
"""Custom Environment that follows gym interface."""
|
||||
|
||||
metadata = {"render_modes": ["human"], "render_fps": 1000/sample_interval}
|
||||
|
||||
def __init__(self, devicenum=0, Ccost=0.2):
|
||||
super().__init__()
|
||||
self.action_space = spaces.Box(low=-1, high=1, shape=(3,), dtype=np.float32)
|
||||
self.observation_space = spaces.Box(low=-5, high=5, shape=(6,), dtype=np.float32)
|
||||
self.fifo_rewards = deque(maxlen=50)
|
||||
lbm.SetDevice(devicenum)
|
||||
lbm.InitAll()
|
||||
lbm.CoreSolver.argtypes = (ctypes.c_int,ctypes.c_float,ctypes.c_float,ctypes.c_float,ctypes.c_float)
|
||||
lbm.CoreSolver.restype = ctypes.POINTER(ctypes.c_float)
|
||||
self.temps_init = lbm.CoreSolver(100*1000, 0.0, 0.0, 0.0, 0.0)
|
||||
self.s = np.array([0.0] * S_DIM, dtype=np.float32)
|
||||
for i in range(S_DIM):
|
||||
self.s[i] = self.temps_init[i]
|
||||
lbm.InitCPUMemory()
|
||||
self.max_steps = max_steps
|
||||
self.current_step = 0
|
||||
self.Ccost = Ccost
|
||||
|
||||
def step(self, action):
|
||||
assert self.action_space.contains(action), "%r (%s) invalid"%(action, type(action))
|
||||
lbm.CoreSolver.argtypes = (ctypes.c_int,ctypes.c_float,ctypes.c_float,ctypes.c_float,ctypes.c_float)
|
||||
lbm.CoreSolver.restype = ctypes.POINTER(ctypes.c_float)
|
||||
action = action_amp * action
|
||||
temps = lbm.CoreSolver(sample_interval, action_weight, action[0], action[1], action[2])
|
||||
for i in range(S_DIM):
|
||||
self.s[i] = temps[i]
|
||||
observation = np.hstack(self.s)
|
||||
cd = self.s[0]+self.s[2]+self.s[4]
|
||||
cl = self.s[1]+self.s[3]+self.s[5]
|
||||
reward = float((1-self.Ccost)*np.exp(-np.abs(cd)/3)+self.Ccost*np.exp(-np.abs(cl)/3))
|
||||
self.fifo_rewards.append(reward)
|
||||
terminated = bool(np.mean(self.fifo_rewards) > 0.9)
|
||||
truncated = bool(np.any(self.s > 3) or np.any(self.s < -3))
|
||||
self.current_step += 1
|
||||
if self.current_step >= self.max_steps:
|
||||
terminated = True
|
||||
info = {}
|
||||
return observation, reward, terminated, truncated, info
|
||||
|
||||
def reset(self, seed=None, Ccost=0.2):
|
||||
lbm.ResetAll()
|
||||
for i in range(S_DIM):
|
||||
self.s[i] = self.temps_init[i]
|
||||
observation = np.hstack(self.s)
|
||||
info = {}
|
||||
self.current_step = 0
|
||||
self.Ccost = Ccost
|
||||
return observation, info
|
||||
|
||||
def render(self, episode=0, numstep=0):
|
||||
lbm.OutputFlow.argtypes = (ctypes.c_int, ctypes.c_int, ctypes.c_int)
|
||||
lbm.OutputFlow(episode, numstep, sample_interval)
|
||||
|
||||
def close(self):
|
||||
lbm.Finalize()
|
||||
@ -1,148 +0,0 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 1,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Using cuda device\n",
|
||||
"Logging to ./tensorboard/PPO_1\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Process ForkServerProcess-2:\n",
|
||||
"Process ForkServerProcess-1:\n",
|
||||
"Traceback (most recent call last):\n",
|
||||
" File \"/home/frank14f/anaconda3/envs/pycuda_3_10/lib/python3.10/multiprocessing/process.py\", line 314, in _bootstrap\n",
|
||||
" self.run()\n",
|
||||
" File \"/home/frank14f/anaconda3/envs/pycuda_3_10/lib/python3.10/multiprocessing/process.py\", line 108, in run\n",
|
||||
" self._target(*self._args, **self._kwargs)\n",
|
||||
" File \"/home/frank14f/anaconda3/envs/pycuda_3_10/lib/python3.10/site-packages/stable_baselines3/common/vec_env/subproc_vec_env.py\", line 35, in _worker\n",
|
||||
" observation, reward, terminated, truncated, info = env.step(data)\n",
|
||||
"ValueError: not enough values to unpack (expected 5, got 4)\n",
|
||||
"Traceback (most recent call last):\n",
|
||||
" File \"/home/frank14f/anaconda3/envs/pycuda_3_10/lib/python3.10/multiprocessing/process.py\", line 314, in _bootstrap\n",
|
||||
" self.run()\n",
|
||||
" File \"/home/frank14f/anaconda3/envs/pycuda_3_10/lib/python3.10/multiprocessing/process.py\", line 108, in run\n",
|
||||
" self._target(*self._args, **self._kwargs)\n",
|
||||
" File \"/home/frank14f/anaconda3/envs/pycuda_3_10/lib/python3.10/site-packages/stable_baselines3/common/vec_env/subproc_vec_env.py\", line 35, in _worker\n",
|
||||
" observation, reward, terminated, truncated, info = env.step(data)\n",
|
||||
"ValueError: not enough values to unpack (expected 5, got 4)\n",
|
||||
"Process ForkServerProcess-4:\n",
|
||||
"Traceback (most recent call last):\n",
|
||||
" File \"/home/frank14f/anaconda3/envs/pycuda_3_10/lib/python3.10/multiprocessing/process.py\", line 314, in _bootstrap\n",
|
||||
" self.run()\n",
|
||||
" File \"/home/frank14f/anaconda3/envs/pycuda_3_10/lib/python3.10/multiprocessing/process.py\", line 108, in run\n",
|
||||
" self._target(*self._args, **self._kwargs)\n",
|
||||
" File \"/home/frank14f/anaconda3/envs/pycuda_3_10/lib/python3.10/site-packages/stable_baselines3/common/vec_env/subproc_vec_env.py\", line 35, in _worker\n",
|
||||
" observation, reward, terminated, truncated, info = env.step(data)\n",
|
||||
"ValueError: not enough values to unpack (expected 5, got 4)\n",
|
||||
"Process ForkServerProcess-3:\n",
|
||||
"Traceback (most recent call last):\n",
|
||||
" File \"/home/frank14f/anaconda3/envs/pycuda_3_10/lib/python3.10/multiprocessing/process.py\", line 314, in _bootstrap\n",
|
||||
" self.run()\n",
|
||||
" File \"/home/frank14f/anaconda3/envs/pycuda_3_10/lib/python3.10/multiprocessing/process.py\", line 108, in run\n",
|
||||
" self._target(*self._args, **self._kwargs)\n",
|
||||
" File \"/home/frank14f/anaconda3/envs/pycuda_3_10/lib/python3.10/site-packages/stable_baselines3/common/vec_env/subproc_vec_env.py\", line 35, in _worker\n",
|
||||
" observation, reward, terminated, truncated, info = env.step(data)\n",
|
||||
"ValueError: not enough values to unpack (expected 5, got 4)\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"ename": "EOFError",
|
||||
"evalue": "",
|
||||
"output_type": "error",
|
||||
"traceback": [
|
||||
"\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
|
||||
"\u001b[0;31mEOFError\u001b[0m Traceback (most recent call last)",
|
||||
"Cell \u001b[0;32mIn[1], line 24\u001b[0m\n\u001b[1;32m 16\u001b[0m vec_env \u001b[38;5;241m=\u001b[39m SubprocVecEnv(env_fns)\n\u001b[1;32m 18\u001b[0m model \u001b[38;5;241m=\u001b[39m PPO(\n\u001b[1;32m 19\u001b[0m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mMlpPolicy\u001b[39m\u001b[38;5;124m\"\u001b[39m, \n\u001b[1;32m 20\u001b[0m env\u001b[38;5;241m=\u001b[39mvec_env, \n\u001b[1;32m 21\u001b[0m n_steps\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m64\u001b[39m,\n\u001b[1;32m 22\u001b[0m tensorboard_log\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m./tensorboard/\u001b[39m\u001b[38;5;124m\"\u001b[39m, \n\u001b[1;32m 23\u001b[0m verbose\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1\u001b[39m)\n\u001b[0;32m---> 24\u001b[0m \u001b[43mmodel\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlearn\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtotal_timesteps\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m128\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m1000\u001b[39;49m\u001b[43m)\u001b[49m\n",
|
||||
"File \u001b[0;32m~/anaconda3/envs/pycuda_3_10/lib/python3.10/site-packages/stable_baselines3/ppo/ppo.py:315\u001b[0m, in \u001b[0;36mPPO.learn\u001b[0;34m(self, total_timesteps, callback, log_interval, tb_log_name, reset_num_timesteps, progress_bar)\u001b[0m\n\u001b[1;32m 306\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mlearn\u001b[39m(\n\u001b[1;32m 307\u001b[0m \u001b[38;5;28mself\u001b[39m: SelfPPO,\n\u001b[1;32m 308\u001b[0m total_timesteps: \u001b[38;5;28mint\u001b[39m,\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 313\u001b[0m progress_bar: \u001b[38;5;28mbool\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m,\n\u001b[1;32m 314\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m SelfPPO:\n\u001b[0;32m--> 315\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlearn\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 316\u001b[0m \u001b[43m \u001b[49m\u001b[43mtotal_timesteps\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtotal_timesteps\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 317\u001b[0m \u001b[43m \u001b[49m\u001b[43mcallback\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcallback\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 318\u001b[0m \u001b[43m \u001b[49m\u001b[43mlog_interval\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mlog_interval\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 319\u001b[0m \u001b[43m \u001b[49m\u001b[43mtb_log_name\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtb_log_name\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 320\u001b[0m \u001b[43m \u001b[49m\u001b[43mreset_num_timesteps\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mreset_num_timesteps\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 321\u001b[0m \u001b[43m \u001b[49m\u001b[43mprogress_bar\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mprogress_bar\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 322\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\n",
|
||||
"File \u001b[0;32m~/anaconda3/envs/pycuda_3_10/lib/python3.10/site-packages/stable_baselines3/common/on_policy_algorithm.py:277\u001b[0m, in \u001b[0;36mOnPolicyAlgorithm.learn\u001b[0;34m(self, total_timesteps, callback, log_interval, tb_log_name, reset_num_timesteps, progress_bar)\u001b[0m\n\u001b[1;32m 274\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39menv \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m 276\u001b[0m \u001b[38;5;28;01mwhile\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mnum_timesteps \u001b[38;5;241m<\u001b[39m total_timesteps:\n\u001b[0;32m--> 277\u001b[0m continue_training \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcollect_rollouts\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43menv\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcallback\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrollout_buffer\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mn_rollout_steps\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mn_steps\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 279\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m continue_training:\n\u001b[1;32m 280\u001b[0m \u001b[38;5;28;01mbreak\u001b[39;00m\n",
|
||||
"File \u001b[0;32m~/anaconda3/envs/pycuda_3_10/lib/python3.10/site-packages/stable_baselines3/common/on_policy_algorithm.py:194\u001b[0m, in \u001b[0;36mOnPolicyAlgorithm.collect_rollouts\u001b[0;34m(self, env, callback, rollout_buffer, n_rollout_steps)\u001b[0m\n\u001b[1;32m 189\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m 190\u001b[0m \u001b[38;5;66;03m# Otherwise, clip the actions to avoid out of bound error\u001b[39;00m\n\u001b[1;32m 191\u001b[0m \u001b[38;5;66;03m# as we are sampling from an unbounded Gaussian distribution\u001b[39;00m\n\u001b[1;32m 192\u001b[0m clipped_actions \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39mclip(actions, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39maction_space\u001b[38;5;241m.\u001b[39mlow, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39maction_space\u001b[38;5;241m.\u001b[39mhigh)\n\u001b[0;32m--> 194\u001b[0m new_obs, rewards, dones, infos \u001b[38;5;241m=\u001b[39m \u001b[43menv\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mstep\u001b[49m\u001b[43m(\u001b[49m\u001b[43mclipped_actions\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 196\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mnum_timesteps \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m env\u001b[38;5;241m.\u001b[39mnum_envs\n\u001b[1;32m 198\u001b[0m \u001b[38;5;66;03m# Give access to local variables\u001b[39;00m\n",
|
||||
"File \u001b[0;32m~/anaconda3/envs/pycuda_3_10/lib/python3.10/site-packages/stable_baselines3/common/vec_env/base_vec_env.py:206\u001b[0m, in \u001b[0;36mVecEnv.step\u001b[0;34m(self, actions)\u001b[0m\n\u001b[1;32m 199\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m 200\u001b[0m \u001b[38;5;124;03mStep the environments with the given action\u001b[39;00m\n\u001b[1;32m 201\u001b[0m \n\u001b[1;32m 202\u001b[0m \u001b[38;5;124;03m:param actions: the action\u001b[39;00m\n\u001b[1;32m 203\u001b[0m \u001b[38;5;124;03m:return: observation, reward, done, information\u001b[39;00m\n\u001b[1;32m 204\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m 205\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mstep_async(actions)\n\u001b[0;32m--> 206\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mstep_wait\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
|
||||
"File \u001b[0;32m~/anaconda3/envs/pycuda_3_10/lib/python3.10/site-packages/stable_baselines3/common/vec_env/subproc_vec_env.py:129\u001b[0m, in \u001b[0;36mSubprocVecEnv.step_wait\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 128\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mstep_wait\u001b[39m(\u001b[38;5;28mself\u001b[39m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m VecEnvStepReturn:\n\u001b[0;32m--> 129\u001b[0m results \u001b[38;5;241m=\u001b[39m [remote\u001b[38;5;241m.\u001b[39mrecv() \u001b[38;5;28;01mfor\u001b[39;00m remote \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mremotes]\n\u001b[1;32m 130\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mwaiting \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m\n\u001b[1;32m 131\u001b[0m obs, rews, dones, infos, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mreset_infos \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mzip\u001b[39m(\u001b[38;5;241m*\u001b[39mresults) \u001b[38;5;66;03m# type: ignore[assignment]\u001b[39;00m\n",
|
||||
"File \u001b[0;32m~/anaconda3/envs/pycuda_3_10/lib/python3.10/site-packages/stable_baselines3/common/vec_env/subproc_vec_env.py:129\u001b[0m, in \u001b[0;36m<listcomp>\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m 128\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mstep_wait\u001b[39m(\u001b[38;5;28mself\u001b[39m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m VecEnvStepReturn:\n\u001b[0;32m--> 129\u001b[0m results \u001b[38;5;241m=\u001b[39m [\u001b[43mremote\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrecv\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m \u001b[38;5;28;01mfor\u001b[39;00m remote \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mremotes]\n\u001b[1;32m 130\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mwaiting \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m\n\u001b[1;32m 131\u001b[0m obs, rews, dones, infos, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mreset_infos \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mzip\u001b[39m(\u001b[38;5;241m*\u001b[39mresults) \u001b[38;5;66;03m# type: ignore[assignment]\u001b[39;00m\n",
|
||||
"File \u001b[0;32m~/anaconda3/envs/pycuda_3_10/lib/python3.10/multiprocessing/connection.py:250\u001b[0m, in \u001b[0;36m_ConnectionBase.recv\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 248\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_check_closed()\n\u001b[1;32m 249\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_check_readable()\n\u001b[0;32m--> 250\u001b[0m buf \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_recv_bytes\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 251\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m _ForkingPickler\u001b[38;5;241m.\u001b[39mloads(buf\u001b[38;5;241m.\u001b[39mgetbuffer())\n",
|
||||
"File \u001b[0;32m~/anaconda3/envs/pycuda_3_10/lib/python3.10/multiprocessing/connection.py:414\u001b[0m, in \u001b[0;36mConnection._recv_bytes\u001b[0;34m(self, maxsize)\u001b[0m\n\u001b[1;32m 413\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_recv_bytes\u001b[39m(\u001b[38;5;28mself\u001b[39m, maxsize\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m):\n\u001b[0;32m--> 414\u001b[0m buf \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_recv\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m4\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m 415\u001b[0m size, \u001b[38;5;241m=\u001b[39m struct\u001b[38;5;241m.\u001b[39munpack(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m!i\u001b[39m\u001b[38;5;124m\"\u001b[39m, buf\u001b[38;5;241m.\u001b[39mgetvalue())\n\u001b[1;32m 416\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m size \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m:\n",
|
||||
"File \u001b[0;32m~/anaconda3/envs/pycuda_3_10/lib/python3.10/multiprocessing/connection.py:383\u001b[0m, in \u001b[0;36mConnection._recv\u001b[0;34m(self, size, read)\u001b[0m\n\u001b[1;32m 381\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m n \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m0\u001b[39m:\n\u001b[1;32m 382\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m remaining \u001b[38;5;241m==\u001b[39m size:\n\u001b[0;32m--> 383\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mEOFError\u001b[39;00m\n\u001b[1;32m 384\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m 385\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mOSError\u001b[39;00m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mgot end of file during message\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
|
||||
"\u001b[0;31mEOFError\u001b[0m: "
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"import os\n",
|
||||
"os.environ['MKL_THREADING_LAYER'] = 'GNU'\n",
|
||||
"import numpy as np\n",
|
||||
"import gymnasium as gym\n",
|
||||
"from env_pinball import CustomEnv\n",
|
||||
"from stable_baselines3 import PPO\n",
|
||||
"from stable_baselines3.common.vec_env import SubprocVecEnv\n",
|
||||
"\n",
|
||||
"def make_env(gpu_id):\n",
|
||||
" def _init():\n",
|
||||
" os.environ[\"CUDA_VISIBLE_DEVICES\"] = str(gpu_id)\n",
|
||||
" return CustomEnv(devicenum=gpu_id)\n",
|
||||
" return _init\n",
|
||||
"\n",
|
||||
"env_fns = [make_env(i) for i in range(4)]\n",
|
||||
"vec_env = SubprocVecEnv(env_fns)\n",
|
||||
"\n",
|
||||
"model = PPO(\n",
|
||||
" \"MlpPolicy\", \n",
|
||||
" env=vec_env, \n",
|
||||
" n_steps=64,\n",
|
||||
" tensorboard_log=\"./tensorboard/\", \n",
|
||||
" verbose=1)\n",
|
||||
"model.learn(total_timesteps=64*1000)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"vec_env = model.get_env()\n",
|
||||
"obs = vec_env.reset()\n",
|
||||
"\n",
|
||||
"n_steps = 0\n",
|
||||
"list_reward = {}\n",
|
||||
"terminated = False\n",
|
||||
"truncated = False\n",
|
||||
"while n_steps < 500 and not terminated and not truncated:\n",
|
||||
" n_steps += 1\n",
|
||||
" action, _states = model.predict(observation=obs)\n",
|
||||
" obs, rewards, dones, info = vec_env.step(action)\n",
|
||||
" list_reward[n_steps] = rewards"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "pycuda_3_10",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 3
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.10.13"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 2
|
||||
}
|
||||
BIN
SB3/lbm_sens.so
BIN
SB3/lbm_sens.so
Binary file not shown.
216
SB3/nohup.py
216
SB3/nohup.py
@ -1,216 +0,0 @@
|
||||
# %%
|
||||
#!/usr/bin/env python3
|
||||
from env_pinball import CustomEnv
|
||||
import os
|
||||
import pickle
|
||||
import random
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.distributions.normal import Normal
|
||||
from torch.utils.tensorboard import SummaryWriter
|
||||
|
||||
env = CustomEnv(devicenum=3)
|
||||
writer = SummaryWriter(log_dir='./tensorboard/DRL')
|
||||
|
||||
# %%
|
||||
class Policy_Network(nn.Module):
|
||||
"""Parametrized Policy Network."""
|
||||
|
||||
def __init__(self, obs_space_dims: int, action_space_dims: int):
|
||||
"""Initializes a neural network that estimates the mean and standard deviation
|
||||
of a normal distribution from which an action is sampled from.
|
||||
|
||||
Args:
|
||||
obs_space_dims: Dimension of the observation space
|
||||
action_space_dims: Dimension of the action space
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
hidden_space1 = 256 # Nothing special with 16, feel free to change
|
||||
hidden_space2 = 256 # Nothing special with 32, feel free to change
|
||||
|
||||
# Shared Network
|
||||
self.shared_net = nn.Sequential(
|
||||
nn.Linear(obs_space_dims, hidden_space1),
|
||||
nn.Tanh(),
|
||||
nn.Linear(hidden_space1, hidden_space2),
|
||||
nn.Tanh(),
|
||||
)
|
||||
|
||||
# Policy Mean specific Linear Layer
|
||||
self.policy_mean_net = nn.Sequential(
|
||||
nn.Linear(hidden_space2, action_space_dims)
|
||||
)
|
||||
|
||||
# Policy Std Dev specific Linear Layer
|
||||
self.policy_stddev_net = nn.Sequential(
|
||||
nn.Linear(hidden_space2, action_space_dims)
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Conditioned on the observation, returns the mean and standard deviation
|
||||
of a normal distribution from which an action is sampled from.
|
||||
|
||||
Args:
|
||||
x: Observation from the environment
|
||||
|
||||
Returns:
|
||||
action_means: predicted mean of the normal distribution
|
||||
action_stddevs: predicted standard deviation of the normal distribution
|
||||
"""
|
||||
shared_features = self.shared_net(x.float())
|
||||
|
||||
action_means = self.policy_mean_net(shared_features)
|
||||
action_means = torch.tanh(action_means)
|
||||
|
||||
action_stddevs = torch.log(
|
||||
1 + torch.exp(self.policy_stddev_net(shared_features))
|
||||
)
|
||||
|
||||
return action_means, action_stddevs
|
||||
|
||||
# %%
|
||||
class REINFORCE:
|
||||
"""REINFORCE algorithm."""
|
||||
|
||||
def __init__(self, obs_space_dims: int, action_space_dims: int):
|
||||
"""Initializes an agent that learns a policy via REINFORCE algorithm [1]
|
||||
to solve the task at hand (Inverted Pendulum v4).
|
||||
|
||||
Args:
|
||||
obs_space_dims: Dimension of the observation space
|
||||
action_space_dims: Dimension of the action space
|
||||
"""
|
||||
|
||||
# Hyperparameters
|
||||
self.learning_rate = 1e-4 # Learning rate for policy optimization
|
||||
self.gamma = 0.99 # Discount factor
|
||||
self.eps = 1e-6 # small number for mathematical stability
|
||||
|
||||
self.probs = [] # Stores probability values of the sampled action
|
||||
self.rewards = [] # Stores the corresponding rewards
|
||||
|
||||
self.net = Policy_Network(obs_space_dims, action_space_dims)
|
||||
self.optimizer = torch.optim.AdamW(self.net.parameters(), lr=self.learning_rate)
|
||||
|
||||
def sample_action(self, state: np.ndarray) -> float:
|
||||
"""Returns an action, conditioned on the policy and observation.
|
||||
|
||||
Args:
|
||||
state: Observation from the environment
|
||||
|
||||
Returns:
|
||||
action: Action to be performed
|
||||
"""
|
||||
state = torch.tensor(np.array([state]))
|
||||
action_means, action_stddevs = self.net(state)
|
||||
|
||||
# create a normal distribution from the predicted
|
||||
# mean and standard deviation and sample an action
|
||||
distrib = Normal(action_means[0] + self.eps, action_stddevs[0] + self.eps)
|
||||
action = distrib.sample()
|
||||
prob = distrib.log_prob(action)
|
||||
|
||||
action = torch.tanh(action)
|
||||
action = action.numpy()
|
||||
|
||||
self.probs.append(prob)
|
||||
|
||||
return action
|
||||
|
||||
def update(self):
|
||||
"""Updates the policy network's weights."""
|
||||
running_g = 0
|
||||
gs = []
|
||||
|
||||
# Discounted return (backwards) - [::-1] will return an array in reverse
|
||||
for R in self.rewards[::-1]:
|
||||
running_g = R + self.gamma * running_g
|
||||
gs.insert(0, running_g)
|
||||
|
||||
deltas = torch.tensor(gs)
|
||||
|
||||
loss = 0
|
||||
# minimize -1 * prob * reward obtained
|
||||
for log_prob, delta in zip(self.probs, deltas):
|
||||
loss += log_prob.mean() * delta * (-1)
|
||||
|
||||
# Update the policy network
|
||||
self.optimizer.zero_grad()
|
||||
loss.backward()
|
||||
self.optimizer.step()
|
||||
|
||||
# Empty / zero out all episode-centric/related variables
|
||||
self.probs = []
|
||||
self.rewards = []
|
||||
|
||||
# %%
|
||||
total_num_episodes = int(5e3) # Total number of episodes
|
||||
obs_space_dims = 6
|
||||
action_space_dims = 3
|
||||
rewards_over_seeds = []
|
||||
MAX_REWARD = 0
|
||||
|
||||
# Check if there is a saved state
|
||||
if os.path.exists('saved_state.pkl'):
|
||||
with open('saved_state.pkl', 'rb') as f:
|
||||
i_seed, episode, agent, reward_over_episodes, rewards_over_seeds, MAX_REWARD = pickle.load(f)
|
||||
os.remove('saved_state.pkl') # Remove the saved state
|
||||
else:
|
||||
i_seed = 0
|
||||
episode = 0
|
||||
agent = None
|
||||
reward_over_episodes = None
|
||||
|
||||
for seed in [1][i_seed:]: # Fibonacci seeds
|
||||
# set seed
|
||||
torch.manual_seed(seed)
|
||||
random.seed(seed)
|
||||
np.random.seed(seed)
|
||||
|
||||
# Reinitialize agent every seed
|
||||
if agent is None or reward_over_episodes is None:
|
||||
agent = REINFORCE(obs_space_dims, action_space_dims)
|
||||
reward_over_episodes = []
|
||||
|
||||
while episode < total_num_episodes+1:
|
||||
obs, info = env.reset(Ccost=0.2+episode/total_num_episodes*0.6)
|
||||
steps = 0
|
||||
done = False
|
||||
terminated = False
|
||||
truncated = False
|
||||
reward_over_steps = []
|
||||
while not done:
|
||||
action = agent.sample_action(obs)
|
||||
obs, reward, terminated, truncated, info = env.step(action)
|
||||
agent.rewards.append(reward)
|
||||
reward_over_steps.append(reward)
|
||||
steps += 1
|
||||
done = terminated or truncated
|
||||
|
||||
avg_reward = np.mean(reward_over_steps[-64:])
|
||||
reward_over_episodes.append(np.array([avg_reward], dtype=np.float32))
|
||||
agent.update()
|
||||
|
||||
if episode % 10 == 0:
|
||||
# print("Episode:", episode, "Average Reward:", int(avg_reward))
|
||||
writer.add_scalar('Average Reward', int(avg_reward), episode)
|
||||
|
||||
if avg_reward > MAX_REWARD:
|
||||
MAX_REWARD = avg_reward
|
||||
with open('saved_model_'+str(seed)+'.pkl', 'wb') as f:
|
||||
pickle.dump((episode + 1, agent, reward_over_episodes, MAX_REWARD), f)
|
||||
|
||||
# Save the current state at the end of each episode
|
||||
with open('saved_state.pkl', 'wb') as f:
|
||||
pickle.dump((i_seed, episode + 1, agent, reward_over_episodes, rewards_over_seeds, MAX_REWARD), f)
|
||||
|
||||
episode += 1
|
||||
episode = 0
|
||||
MAX_REWARD = 0
|
||||
i_seed += 1
|
||||
rewards_over_seeds.append(reward_over_episodes)
|
||||
agent = None # Reset the agent
|
||||
reward_over_episodes = None # Reset the reward_over_episodes
|
||||
# %%
|
||||
BIN
SB3/ppo_1.zip
BIN
SB3/ppo_1.zip
Binary file not shown.
BIN
SB3/ppo_2.zip
BIN
SB3/ppo_2.zip
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
9
configs/config_cuda.json
Normal file
9
configs/config_cuda.json
Normal file
@ -0,0 +1,9 @@
|
||||
{
|
||||
"multi_gpu": false,
|
||||
"gpu_connection": "NVLink",
|
||||
"required_cuda_capability": "7.0",
|
||||
"threads_per_block": 128,
|
||||
"X_1U": 128,
|
||||
"Y_1U": 32,
|
||||
"Z_1U": 1
|
||||
}
|
||||
13
configs/config_flowfield.json
Normal file
13
configs/config_flowfield.json
Normal file
@ -0,0 +1,13 @@
|
||||
{
|
||||
"data_type": "FP32",
|
||||
"dimensionality": 2,
|
||||
"lattice": 9,
|
||||
"field_dim_in_U": [10, 16, 1],
|
||||
"viscosity": 0.004,
|
||||
"velocity": 0.01,
|
||||
"boundary_conditions": {
|
||||
"x": ["parabolic", "outflow"],
|
||||
"y": ["noslip", "noslip"],
|
||||
"z": ["none", "none"]
|
||||
}
|
||||
}
|
||||
3
configs/config_gym.json
Normal file
3
configs/config_gym.json
Normal file
@ -0,0 +1,3 @@
|
||||
{
|
||||
|
||||
}
|
||||
@ -1,9 +0,0 @@
|
||||
import pycuda.driver as cuda
|
||||
import pycuda.autoinit
|
||||
from pycuda.compiler import SourceModule
|
||||
import numpy as np
|
||||
|
||||
with open('./cuda/kernel.cu', 'r') as file_k:
|
||||
code = file_k.read()
|
||||
|
||||
kernel = SourceModule(code)
|
||||
BIN
output.png
BIN
output.png
Binary file not shown.
|
Before Width: | Height: | Size: 73 KiB |
BIN
profile.nvvp
Normal file
BIN
profile.nvvp
Normal file
Binary file not shown.
BIN
scripts/__pycache__/gym_env.cpython-310.pyc
Normal file
BIN
scripts/__pycache__/gym_env.cpython-310.pyc
Normal file
Binary file not shown.
57
scripts/d1a3o12.py
Normal file
57
scripts/d1a3o12.py
Normal file
@ -0,0 +1,57 @@
|
||||
import os
|
||||
os.environ['MKL_THREADING_LAYER'] = 'GNU'
|
||||
os.environ["OMP_NUM_THREADS"] = "8"
|
||||
os.environ["MKL_NUM_THREADS"] = "8"
|
||||
import torch
|
||||
import numpy as np
|
||||
from torch.nn import Module
|
||||
import gymnasium as gym
|
||||
from gym_env import CustomEnv
|
||||
from stable_baselines3 import PPO
|
||||
from stable_baselines3.common.vec_env import SubprocVecEnv
|
||||
from stable_baselines3.common.vec_env import DummyVecEnv
|
||||
from sb3_contrib import RecurrentPPO
|
||||
from torch.utils.tensorboard import SummaryWriter
|
||||
|
||||
current_dir = os.path.dirname(os.path.abspath("__file__"))
|
||||
parent_dir = os.path.abspath(os.path.join(current_dir, os.pardir))
|
||||
|
||||
class Sin(Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
def forward(self, x):
|
||||
return torch.sin(x)
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
vec_env = CustomEnv(device_id=1)
|
||||
name = "d1a3o12_c1"
|
||||
|
||||
model = PPO.load(os.path.join(parent_dir, "models", "d1a3o12_a0"), env=vec_env, device=torch.device("cuda:1"))
|
||||
|
||||
# model = PPO(
|
||||
# "MlpPolicy",
|
||||
# policy_kwargs=dict(activation_fn=Sin),
|
||||
# env=vec_env,
|
||||
# device=torch.device("cuda:1"),
|
||||
# verbose=0)
|
||||
|
||||
writer = SummaryWriter(log_dir=os.path.join(parent_dir, "tensorboard", name))
|
||||
max_reward = 0
|
||||
|
||||
for i in range(100):
|
||||
model.learn(total_timesteps=480)
|
||||
test_env = model.get_env()
|
||||
test_obs = test_env.reset()
|
||||
list_reward = []
|
||||
for step in range(480):
|
||||
test_action, _states = model.predict(observation=test_obs)
|
||||
test_obs, test_rewards, test_dones, info = test_env.step(test_action)
|
||||
list_reward.append(test_rewards)
|
||||
|
||||
avg_reward = np.mean(list_reward[-240:])
|
||||
writer.add_scalar('Reward', np.mean(avg_reward), i)
|
||||
if avg_reward > max_reward:
|
||||
max_reward = avg_reward
|
||||
model.save(os.path.join(parent_dir, "models", name + ".zip"))
|
||||
198
scripts/gym_env.py
Normal file
198
scripts/gym_env.py
Normal file
@ -0,0 +1,198 @@
|
||||
import gymnasium as gym
|
||||
import numpy as np
|
||||
from gymnasium import spaces
|
||||
import ctypes
|
||||
from collections import deque
|
||||
from typing import Tuple
|
||||
import sys
|
||||
import os
|
||||
import threading
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from concurrent.futures import ProcessPoolExecutor
|
||||
import queue
|
||||
|
||||
os.environ["OMP_NUM_THREADS"] = "1"
|
||||
os.environ["MKL_NUM_THREADS"] = "1"
|
||||
|
||||
current_dir = os.path.dirname(os.path.abspath("__file__"))
|
||||
parent_dir = os.path.abspath(os.path.join(current_dir, os.pardir))
|
||||
sys.path.append(parent_dir)
|
||||
from CelerisLab import FlowField
|
||||
from CelerisLab import utils
|
||||
|
||||
config_cuda = utils.load_cuda_config(
|
||||
os.path.join(parent_dir, "configs", "config_cuda.json")
|
||||
)
|
||||
config_field = utils.load_flow_field_config(
|
||||
os.path.join(parent_dir, "configs", "config_flowfield.json")
|
||||
)
|
||||
|
||||
S_DIM, A_DIM = 12, 3
|
||||
U0 = config_field.velocity
|
||||
T0 = 1000
|
||||
SAMPLE_INTERVAL = 800
|
||||
FIFO_LEN = 120
|
||||
CONV_LEN = 60
|
||||
MAX_STEPS = 640
|
||||
if config_field.data_type == "FP32":
|
||||
DATA_TYPE = np.float32
|
||||
else:
|
||||
raise ValueError(f"Unsupported data type {config_field.data_type}.")
|
||||
|
||||
|
||||
class CustomEnv(gym.Env):
|
||||
"""Custom Environment that follows gym interface."""
|
||||
|
||||
metadata = {"render_modes": ["human"], "render_fps": T0 / SAMPLE_INTERVAL}
|
||||
|
||||
def __init__(self, device_id=0):
|
||||
super().__init__()
|
||||
self.action_space = spaces.Box(low=-1, high=1, shape=(A_DIM,), dtype=DATA_TYPE)
|
||||
self.observation_space = spaces.Box(
|
||||
low=-1, high=1, shape=(S_DIM,), dtype=DATA_TYPE
|
||||
)
|
||||
self.fifo_states = deque(maxlen=FIFO_LEN)
|
||||
self.target_states = np.empty((0, 6), dtype=DATA_TYPE)
|
||||
self.force_norm_fact = 1.0
|
||||
self.sens_norm_fact = np.ones(6, dtype=DATA_TYPE)
|
||||
self.sens_deviation = np.zeros(6, dtype=DATA_TYPE)
|
||||
|
||||
self.flow_field = FlowField(config_field, config_cuda, device_id)
|
||||
L0 = 20
|
||||
U0 = config_field.velocity
|
||||
NX = self.flow_field.FIELD_SHAPE[0]
|
||||
NY = self.flow_field.FIELD_SHAPE[1]
|
||||
center: Tuple[float, float, float] = (10 * L0, (NY - 1) / 2, 0)
|
||||
self.flow_field.add_cylinder(center, L0)
|
||||
center: Tuple[float, float, float] = (40 * L0, (NY - 1) / 2 + 2 * L0, 0)
|
||||
self.flow_field.add_sensor(center, L0 / 4)
|
||||
center: Tuple[float, float, float] = (40 * L0, (NY - 1) / 2, 0)
|
||||
self.flow_field.add_sensor(center, L0 / 4)
|
||||
center: Tuple[float, float, float] = (40 * L0, (NY - 1) / 2 - 2 * L0, 0)
|
||||
self.flow_field.add_sensor(center, L0 / 4)
|
||||
self.flow_field.run(int(4*NX/U0), np.zeros(4, dtype=DATA_TYPE))
|
||||
|
||||
for i in range(FIFO_LEN):
|
||||
self.flow_field.run(SAMPLE_INTERVAL, np.zeros(4, dtype=DATA_TYPE))
|
||||
new_state = self.flow_field.obs.copy()[2:8]
|
||||
self.target_states = np.vstack((self.target_states, new_state))
|
||||
|
||||
self.flow_field.apply_ddf()
|
||||
center: Tuple[float, float, float] = (30 * L0, (NY - 1) / 2, 0)
|
||||
self.flow_field.add_cylinder(center, L0 / 2)
|
||||
center: Tuple[float, float, float] = (31.3 * L0, (NY - 1) / 2 + 0.75 * L0, 0)
|
||||
self.flow_field.add_cylinder(center, L0 / 2)
|
||||
center: Tuple[float, float, float] = (31.3 * L0, (NY - 1) / 2 - 0.75 * L0, 0)
|
||||
self.flow_field.add_cylinder(center, L0 / 2)
|
||||
self.flow_field.run(int(4*NX/U0), np.zeros(7, dtype=DATA_TYPE))
|
||||
self.flow_field.get_ddf()
|
||||
|
||||
for i in range(FIFO_LEN):
|
||||
self.flow_field.run(SAMPLE_INTERVAL, np.zeros(7, dtype=DATA_TYPE))
|
||||
self.fifo_states.append(self.flow_field.obs.copy()[2:14])
|
||||
|
||||
temp_states = np.array(self.fifo_states)
|
||||
self.force_norm_fact = 6 * np.max(np.abs(temp_states[:, 6:12]))
|
||||
for i in range(6):
|
||||
self.sens_deviation[i] = np.mean(temp_states[:, i])
|
||||
self.sens_norm_fact[i] = 5 * np.max(np.abs(temp_states[:, i] - self.sens_deviation[i]))
|
||||
self.target_states[:, i] = (self.target_states[:, i] - self.sens_deviation[i]) / self.sens_norm_fact[i]
|
||||
|
||||
|
||||
def step(self, action):
|
||||
assert self.action_space.contains(action), "%r (%s) invalid" % (
|
||||
action,
|
||||
type(action),
|
||||
)
|
||||
|
||||
# barrier = threading.Barrier(2)
|
||||
result_queue = queue.Queue()
|
||||
|
||||
def run_flow_field(action):
|
||||
self.flow_field.context.push()
|
||||
U0 = config_field.velocity
|
||||
try:
|
||||
temp = np.zeros(7, dtype=DATA_TYPE)
|
||||
temp[4:7] = np.array((action*8+[0,-4,4])*U0, dtype=DATA_TYPE)
|
||||
self.flow_field.run(SAMPLE_INTERVAL, temp)
|
||||
finally:
|
||||
self.flow_field.context.pop()
|
||||
# barrier.wait()
|
||||
self.fifo_states.append(self.flow_field.obs.copy()[2:14])
|
||||
|
||||
def calc_reward():
|
||||
states = np.array(self.fifo_states)
|
||||
forces = states[-1, 6:12] / self.force_norm_fact
|
||||
cd = (forces[0] + forces[2] + forces[4]) / 3
|
||||
cl = (forces[1] + forces[3] + forces[5]) / 3
|
||||
sens = (states[-1, 0:6] - self.sens_deviation) / self.sens_norm_fact
|
||||
|
||||
similarities = 0.0
|
||||
|
||||
def calc_lag(target, state):
|
||||
target_mean = np.mean(target)
|
||||
state_mean = np.mean(state)
|
||||
|
||||
correlation = np.correlate(target - target_mean, state - state_mean, "full")
|
||||
lags = np.arange(-len(target) + 1, len(target))
|
||||
max_lag = lags[np.argmax(correlation)]
|
||||
return max_lag
|
||||
|
||||
def calc_sim(target, state, lag):
|
||||
target_mean = np.mean(target)
|
||||
state_mean = np.mean(state)
|
||||
target_std = np.std(target)
|
||||
|
||||
aligned_state = np.roll(state, lag)
|
||||
|
||||
if lag >= 0:
|
||||
seq_target = target[-CONV_LEN:]-target_mean
|
||||
seq_state = aligned_state[-CONV_LEN:]-state_mean
|
||||
else:
|
||||
seq_target = target[:CONV_LEN]-target_mean
|
||||
seq_state = aligned_state[:CONV_LEN]-state_mean
|
||||
|
||||
seq_diff = seq_target - seq_state
|
||||
sim_cor = 10*(np.corrcoef(seq_target, seq_state)[0, 1] - 1)
|
||||
sim_div = -np.abs((target_mean - state_mean) / target_std * 0.75)
|
||||
sim_amp = -np.abs(np.std(seq_diff) / target_std * 2)
|
||||
|
||||
return np.exp((sim_cor + sim_div + sim_amp) / 3)
|
||||
|
||||
id_sens = 0
|
||||
target_seq = self.target_states[:, id_sens]
|
||||
state_seq = (states[:, id_sens] - self.sens_deviation[id_sens]) / self.sens_norm_fact[id_sens]
|
||||
lag = calc_lag(target_seq, state_seq)
|
||||
similarities += calc_sim(target_seq, state_seq, lag) / 6
|
||||
|
||||
for i in range(1, 6):
|
||||
target_seq = self.target_states[:, i]
|
||||
state_seq = (states[:, i] - self.sens_deviation[i]) / self.sens_norm_fact[i]
|
||||
similarities += calc_sim(target_seq, state_seq, lag) / 6
|
||||
|
||||
reward_cd = np.exp(-np.abs(cd * 80))
|
||||
reward_cl = np.exp(-np.abs(cl * 20))
|
||||
# reward_sim = np.exp(2 * (similarities - 1))
|
||||
reward_sim = similarities
|
||||
reward = np.minimum(0.3 * reward_cd + 0.3 * reward_cl + 0.4 * reward_sim, 1.0)
|
||||
# barrier.wait()
|
||||
result_queue.put((np.hstack([forces, sens]), reward))
|
||||
|
||||
run_flow_field(action)
|
||||
calc_reward()
|
||||
observation, reward = result_queue.get()
|
||||
|
||||
truncated = bool(np.any(observation > 1) or np.any(observation < -1))
|
||||
observation = np.clip(observation, -1, 1)
|
||||
# truncated = False
|
||||
return observation, float(reward), False, truncated, {}
|
||||
|
||||
def reset(self, seed=None):
|
||||
self.flow_field.apply_ddf()
|
||||
return np.zeros(S_DIM, dtype=np.float32), {}
|
||||
|
||||
def render(self, mode="human"):
|
||||
pass
|
||||
|
||||
def close(self):
|
||||
self.flow_field.__del__()
|
||||
182
scripts/jupyter.ipynb
Normal file
182
scripts/jupyter.ipynb
Normal file
File diff suppressed because one or more lines are too long
0
scripts/nohup.out
Normal file
0
scripts/nohup.out
Normal file
0
scripts/nohup1.out
Normal file
0
scripts/nohup1.out
Normal file
0
scripts/nohup_c.out
Normal file
0
scripts/nohup_c.out
Normal file
318
scripts/test.ipynb
Normal file
318
scripts/test.ipynb
Normal file
File diff suppressed because one or more lines are too long
17
setup.py
Normal file
17
setup.py
Normal file
@ -0,0 +1,17 @@
|
||||
from setuptools import setup, find_packages
|
||||
|
||||
setup(
|
||||
name='CelerisLab',
|
||||
version='0.1',
|
||||
packages=find_packages(),
|
||||
install_requires=[
|
||||
'pycuda',
|
||||
'numpy',
|
||||
'json'
|
||||
],
|
||||
entry_points={
|
||||
'console_scripts': [
|
||||
'CelerisLab=CelerisLab.driver:main',
|
||||
],
|
||||
},
|
||||
)
|
||||
13
test.cu
13
test.cu
@ -1,13 +0,0 @@
|
||||
#include <stdio.h>
|
||||
#include <cuda.h>
|
||||
|
||||
extern "C" {
|
||||
__global__ void add(float *x, float *y)
|
||||
{
|
||||
int index = threadIdx.x;
|
||||
int stride = blockDim.x;
|
||||
|
||||
for (int i = index; i < 100; i += stride)
|
||||
y[i] = x[i] + y[i];
|
||||
}
|
||||
}
|
||||
84
test.ipynb
84
test.ipynb
@ -1,84 +0,0 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": []
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 1,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"ename": "TypeError",
|
||||
"evalue": "No registered converter was able to produce a C++ rvalue of type unsigned int from this Python object of type numpy.int64",
|
||||
"output_type": "error",
|
||||
"traceback": [
|
||||
"\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
|
||||
"\u001b[0;31mTypeError\u001b[0m Traceback (most recent call last)",
|
||||
"Cell \u001b[0;32mIn[1], line 25\u001b[0m\n\u001b[1;32m 22\u001b[0m drv\u001b[38;5;241m.\u001b[39mmemcpy_htod(y_gpu, y)\n\u001b[1;32m 24\u001b[0m \u001b[38;5;66;03m# 调用函数\u001b[39;00m\n\u001b[0;32m---> 25\u001b[0m \u001b[43madd_func\u001b[49m\u001b[43m(\u001b[49m\u001b[43mn\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mx_gpu\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43my_gpu\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mblock\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m256\u001b[39;49m\u001b[43m,\u001b[49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m,\u001b[49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mgrid\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mn\u001b[49m\u001b[38;5;241;43m/\u001b[39;49m\u001b[38;5;241;43m/\u001b[39;49m\u001b[38;5;241;43m256\u001b[39;49m\u001b[43m,\u001b[49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 27\u001b[0m \u001b[38;5;66;03m# 将结果复制回CPU\u001b[39;00m\n\u001b[1;32m 28\u001b[0m drv\u001b[38;5;241m.\u001b[39mmemcpy_dtoh(y, y_gpu)\n",
|
||||
"File \u001b[0;32m~/anaconda3/envs/pycuda_3_10/lib/python3.10/site-packages/pycuda/driver.py:502\u001b[0m, in \u001b[0;36m_add_functionality.<locals>.function_call\u001b[0;34m(func, *args, **kwargs)\u001b[0m\n\u001b[1;32m 498\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtime\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m time\n\u001b[1;32m 500\u001b[0m start_time \u001b[38;5;241m=\u001b[39m time()\n\u001b[0;32m--> 502\u001b[0m \u001b[43mfunc\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_launch_kernel\u001b[49m\u001b[43m(\u001b[49m\u001b[43mgrid\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mblock\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43marg_buf\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mshared\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m)\u001b[49m\n\u001b[1;32m 504\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m post_handlers \u001b[38;5;129;01mor\u001b[39;00m time_kernel:\n\u001b[1;32m 505\u001b[0m Context\u001b[38;5;241m.\u001b[39msynchronize()\n",
|
||||
"\u001b[0;31mTypeError\u001b[0m: No registered converter was able to produce a C++ rvalue of type unsigned int from this Python object of type numpy.int64"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"import pycuda.autoinit\n",
|
||||
"import pycuda.driver as drv\n",
|
||||
"import numpy as np\n",
|
||||
"\n",
|
||||
"# 加载PTX文件\n",
|
||||
"mod = drv.module_from_file(\"test.ptx\")\n",
|
||||
"\n",
|
||||
"# 获取函数\n",
|
||||
"add_func = mod.get_function(\"add\")\n",
|
||||
"\n",
|
||||
"# 创建数据\n",
|
||||
"n = np.uint32(100) # Convert to unsigned int\n",
|
||||
"x = np.random.rand(n).astype(np.float32)\n",
|
||||
"y = np.random.rand(n).astype(np.float32)\n",
|
||||
"\n",
|
||||
"# 分配内存\n",
|
||||
"x_gpu = drv.mem_alloc(x.nbytes)\n",
|
||||
"y_gpu = drv.mem_alloc(y.nbytes)\n",
|
||||
"\n",
|
||||
"# 将数据复制到GPU\n",
|
||||
"drv.memcpy_htod(x_gpu, x)\n",
|
||||
"drv.memcpy_htod(y_gpu, y)\n",
|
||||
"\n",
|
||||
"# 调用函数\n",
|
||||
"add_func(x_gpu, y_gpu, block=(256,1,1), grid=(n//256,1))\n",
|
||||
"\n",
|
||||
"# 将结果复制回CPU\n",
|
||||
"drv.memcpy_dtoh(y, y_gpu)\n",
|
||||
"\n",
|
||||
"# 检查结果\n",
|
||||
"print(y)"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "pycuda_3_10",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 3
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.10.13"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 2
|
||||
}
|
||||
104
test.ptx
104
test.ptx
@ -1,104 +0,0 @@
|
||||
//
|
||||
// Generated by NVIDIA NVVM Compiler
|
||||
//
|
||||
// Compiler Build ID: CL-32688072
|
||||
// Cuda compilation tools, release 12.1, V12.1.105
|
||||
// Based on NVVM 7.0.1
|
||||
//
|
||||
|
||||
.version 8.1
|
||||
.target sm_52
|
||||
.address_size 64
|
||||
|
||||
// .globl add
|
||||
|
||||
.visible .entry add(
|
||||
.param .u64 add_param_0,
|
||||
.param .u64 add_param_1
|
||||
)
|
||||
{
|
||||
.reg .pred %p<6>;
|
||||
.reg .f32 %f<16>;
|
||||
.reg .b32 %r<22>;
|
||||
.reg .b64 %rd<25>;
|
||||
|
||||
|
||||
ld.param.u64 %rd11, [add_param_0];
|
||||
ld.param.u64 %rd12, [add_param_1];
|
||||
cvta.to.global.u64 %rd1, %rd12;
|
||||
cvta.to.global.u64 %rd2, %rd11;
|
||||
mov.u32 %r1, %ntid.x;
|
||||
mov.u32 %r20, %tid.x;
|
||||
setp.gt.s32 %p1, %r20, 99;
|
||||
@%p1 bra $L__BB0_7;
|
||||
|
||||
mov.u32 %r12, 99;
|
||||
sub.s32 %r13, %r12, %r20;
|
||||
div.u32 %r3, %r13, %r1;
|
||||
add.s32 %r14, %r3, 1;
|
||||
and.b32 %r19, %r14, 3;
|
||||
setp.eq.s32 %p2, %r19, 0;
|
||||
@%p2 bra $L__BB0_4;
|
||||
|
||||
mul.wide.s32 %rd13, %r20, 4;
|
||||
add.s64 %rd24, %rd1, %rd13;
|
||||
mul.wide.s32 %rd4, %r1, 4;
|
||||
add.s64 %rd23, %rd2, %rd13;
|
||||
|
||||
$L__BB0_3:
|
||||
.pragma "nounroll";
|
||||
ld.global.f32 %f1, [%rd24];
|
||||
ld.global.f32 %f2, [%rd23];
|
||||
add.f32 %f3, %f2, %f1;
|
||||
st.global.f32 [%rd24], %f3;
|
||||
add.s32 %r20, %r20, %r1;
|
||||
add.s64 %rd24, %rd24, %rd4;
|
||||
add.s64 %rd23, %rd23, %rd4;
|
||||
add.s32 %r19, %r19, -1;
|
||||
setp.ne.s32 %p3, %r19, 0;
|
||||
@%p3 bra $L__BB0_3;
|
||||
|
||||
$L__BB0_4:
|
||||
setp.lt.u32 %p4, %r3, 3;
|
||||
@%p4 bra $L__BB0_7;
|
||||
|
||||
mul.wide.s32 %rd10, %r1, 4;
|
||||
|
||||
$L__BB0_6:
|
||||
mul.wide.s32 %rd14, %r20, 4;
|
||||
add.s64 %rd15, %rd2, %rd14;
|
||||
add.s64 %rd16, %rd1, %rd14;
|
||||
ld.global.f32 %f4, [%rd16];
|
||||
ld.global.f32 %f5, [%rd15];
|
||||
add.f32 %f6, %f5, %f4;
|
||||
st.global.f32 [%rd16], %f6;
|
||||
add.s64 %rd17, %rd15, %rd10;
|
||||
add.s64 %rd18, %rd16, %rd10;
|
||||
ld.global.f32 %f7, [%rd18];
|
||||
ld.global.f32 %f8, [%rd17];
|
||||
add.f32 %f9, %f8, %f7;
|
||||
st.global.f32 [%rd18], %f9;
|
||||
add.s32 %r15, %r20, %r1;
|
||||
add.s32 %r16, %r15, %r1;
|
||||
add.s64 %rd19, %rd17, %rd10;
|
||||
add.s64 %rd20, %rd18, %rd10;
|
||||
ld.global.f32 %f10, [%rd20];
|
||||
ld.global.f32 %f11, [%rd19];
|
||||
add.f32 %f12, %f11, %f10;
|
||||
st.global.f32 [%rd20], %f12;
|
||||
add.s32 %r17, %r16, %r1;
|
||||
add.s64 %rd21, %rd19, %rd10;
|
||||
add.s64 %rd22, %rd20, %rd10;
|
||||
ld.global.f32 %f13, [%rd22];
|
||||
ld.global.f32 %f14, [%rd21];
|
||||
add.f32 %f15, %f14, %f13;
|
||||
st.global.f32 [%rd22], %f15;
|
||||
add.s32 %r20, %r17, %r1;
|
||||
setp.lt.s32 %p5, %r20, 100;
|
||||
@%p5 bra $L__BB0_6;
|
||||
|
||||
$L__BB0_7:
|
||||
ret;
|
||||
|
||||
}
|
||||
|
||||
Loading…
Reference in New Issue
Block a user