87 lines
3.1 KiB
Python
87 lines
3.1 KiB
Python
# 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) |