62 lines
1.9 KiB
Python
62 lines
1.9 KiB
Python
# CelerisLab/cuda/obs_layout.py
|
|
"""Packed ``obs`` buffer layout for force / torque / sensor telemetry.
|
|
|
|
Usage::
|
|
Host code (e.g. ``ObjectManager``) and ``compiler_v2.generate_config`` must use
|
|
the same ``dim`` and ``n_objects`` so allocation sizes match device macros.
|
|
|
|
Design::
|
|
Segments are contiguous in this order:
|
|
1) force : ``max(n_objects, 1) * dim`` floats
|
|
2) torque : ``max(n_objects, 1) * (1 if dim==2 else dim)`` floats
|
|
3) sensor : ``max(n_objects, 1) * dim`` floats
|
|
"""
|
|
|
|
from typing import NamedTuple
|
|
|
|
|
|
class ObsLayout(NamedTuple):
|
|
"""Integer offsets and sizes in **float** elements (not bytes)."""
|
|
|
|
n_slots: int
|
|
torque_components: int
|
|
slot_stride_floats: int
|
|
force_floats: int
|
|
torque_floats: int
|
|
force0_floats: int
|
|
torque0_floats: int
|
|
sensor0_floats: int
|
|
total_floats: int
|
|
|
|
|
|
def obs_layout(dim: int, n_objects: int) -> ObsLayout:
|
|
"""Compute packed ``obs`` indices for force, torque, and sensor segments.
|
|
|
|
Args:
|
|
dim: Spatial dimension (2 or 3); must match ``LBMConfig.dim``.
|
|
n_objects: Number of registered objects; may be zero.
|
|
|
|
Returns:
|
|
:class:`ObsLayout` with segment sizes and float offsets.
|
|
"""
|
|
n_slots = max(int(n_objects), 1)
|
|
torque_components = 1 if int(dim) == 2 else int(dim)
|
|
slot_stride_floats = n_slots * dim
|
|
force_floats = slot_stride_floats
|
|
torque_floats = n_slots * torque_components
|
|
force0_floats = 0
|
|
torque0_floats = force0_floats + force_floats
|
|
sensor0_floats = torque0_floats + torque_floats
|
|
total_floats = sensor0_floats + slot_stride_floats
|
|
return ObsLayout(
|
|
n_slots=n_slots,
|
|
torque_components=torque_components,
|
|
slot_stride_floats=slot_stride_floats,
|
|
force_floats=force_floats,
|
|
torque_floats=torque_floats,
|
|
force0_floats=force0_floats,
|
|
torque0_floats=torque0_floats,
|
|
sensor0_floats=sensor0_floats,
|
|
total_floats=total_floats,
|
|
)
|