257 lines
8.2 KiB
Python
257 lines
8.2 KiB
Python
# 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
|
|
)
|