215 lines
9.7 KiB
Python
215 lines
9.7 KiB
Python
"""Conservative state dependencies for compiler-owned expressions.
|
|
|
|
This is structural bookkeeping, not C parsing or numerical Jacobian evaluation.
|
|
Native multi-output operations, projections and mode-dependent writes must be
|
|
registered explicitly by the lowering path. Unresolved reachable inputs disable
|
|
coloring for the complete model.
|
|
"""
|
|
from collections import deque
|
|
from dataclasses import dataclass
|
|
import hashlib
|
|
import json
|
|
import re
|
|
|
|
|
|
# Reviewed expression leaves. Unknown arrays/scalars remain unresolved instead
|
|
# of silently becoming constants; the flow schedule supplies its own explicit IR.
|
|
_ARRAY = re.compile(r'\b[A-Za-z_]\w*\[\d+\](?:\.[A-Za-z_]\w*)?|\bgas_\d+\.[A-Za-z_]\w*')
|
|
_IDENTIFIER = re.compile(r'\b[A-Za-z_]\w*\b')
|
|
_NUMBER = re.compile(r'(?<![\w.])(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][+-]?\d+)?')
|
|
_FUNCTIONS = frozenset({
|
|
'fmax', 'fmin', 'fabs', 'sqrt', 'copysign', 'pow',
|
|
'native_signal', 'native_contact', 'native_limit_force', 'native_dry_friction',
|
|
'native_pipe_flow_context', 'native_pipe_flow_cached_context',
|
|
'native_temperature_ph_context', 'native_density_context',
|
|
})
|
|
_CONSTANTS = frozenset({'t', 'properties', 'NAN', 'INFINITY', 'NULL', 'true', 'false'})
|
|
|
|
|
|
def expression_inputs(expression):
|
|
"""Extract leaves from a reviewed expression; unfamiliar names fail closed."""
|
|
inputs = set()
|
|
def leaf(match):
|
|
key = match.group()
|
|
# Per-evaluation cache scratch is an implementation detail of reviewed
|
|
# kernels; each value is determined by that kernel's explicit arguments.
|
|
if not re.fullmatch(r'pipe_cache\[\d+\]', key):
|
|
inputs.add(key)
|
|
return ' '
|
|
remainder = _ARRAY.sub(leaf, expression)
|
|
remainder = _NUMBER.sub(' ', remainder)
|
|
for name in _IDENTIFIER.findall(remainder):
|
|
if name in _FUNCTIONS or name in _CONSTANTS or re.fullmatch(r'(?:medium|signal)_\d+', name):
|
|
continue
|
|
inputs.add(name)
|
|
return frozenset(inputs)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class JacobianStructure:
|
|
state_count: int
|
|
rows: tuple[tuple[int, ...], ...] = ()
|
|
colors: tuple[int, ...] = ()
|
|
reason: str | None = None
|
|
|
|
@classmethod
|
|
def dense(cls, state_count, reason):
|
|
return cls(state_count, reason=reason)
|
|
|
|
@property
|
|
def enabled(self):
|
|
return self.reason is None
|
|
|
|
@property
|
|
def color_count(self):
|
|
return max(self.colors, default=-1) + 1 if self.enabled else self.state_count
|
|
|
|
@property
|
|
def nonzeros(self):
|
|
return sum(map(len, self.rows)) if self.enabled else self.state_count ** 2
|
|
|
|
def manifest(self, *, canonical_rhs=False):
|
|
# Use proven coloring automatically when it reduces finite-difference
|
|
# groups; unsupported or unhelpful patterns retain dense differences.
|
|
runtime_eligible = self.enabled and 0 < self.color_count < self.state_count
|
|
fallback_reason = None if runtime_eligible else (
|
|
self.reason or 'Coloring does not reduce finite-difference groups')
|
|
pattern = {'rows': self.rows, 'columnColors': self.colors}
|
|
policy = ('CVODE colored forward differences; canonical-property-cache RHS' if canonical_rhs else 'CVODE colored forward differences') if runtime_eligible else 'CVODE default dense differences'
|
|
return {
|
|
'policyScope': 'default-runtime',
|
|
'defaultRuntimePolicy': policy,
|
|
'verification': '--verify-jacobian',
|
|
'policy': policy,
|
|
'rhsPolicy': 'canonical-property-cache finite differences' if canonical_rhs and runtime_eligible else 'ordinary model_eval',
|
|
'runtimeEligible': runtime_eligible, 'runtimeFallbackReason': fallback_reason,
|
|
'canonicalRhs': canonical_rhs, 'ordinaryRhsUnchanged': True,
|
|
'enabled': self.enabled, 'reason': self.reason, 'stateCount': self.state_count,
|
|
'nonzeros': self.nonzeros, 'density': self.nonzeros / self.state_count ** 2 if self.state_count else 0,
|
|
'colorCount': self.color_count,
|
|
'patternSha256': hashlib.sha256(json.dumps(pattern, separators=(',', ':')).encode()).hexdigest() if self.enabled else None,
|
|
'columnColors': list(self.colors) if self.enabled else None,
|
|
}
|
|
|
|
def header_lines(self, *, canonical_rhs=False):
|
|
lines = [f'#define MODEL_JACOBIAN_COLORED {int(self.enabled)}',
|
|
f'#define MODEL_JACOBIAN_CANONICAL_RHS {int(canonical_rhs)}',
|
|
f'#define MODEL_JACOBIAN_COLOR_COUNT {self.color_count}',
|
|
f'#define MODEL_JACOBIAN_NNZ {self.nonzeros}']
|
|
if canonical_rhs:
|
|
lines += ['int model_eval_jacobian(double t, const double *y, double *dy, double *w);']
|
|
if self.enabled:
|
|
lines += ['extern const int model_jacobian_column_color[NSTATES];',
|
|
'extern const int model_jacobian_col_ptr[NSTATES+1];',
|
|
'extern const int model_jacobian_row_index[MODEL_JACOBIAN_NNZ];']
|
|
return lines
|
|
|
|
def source_lines(self):
|
|
if not self.enabled:
|
|
return []
|
|
columns = [[] for _ in range(self.state_count)]
|
|
for row, entries in enumerate(self.rows):
|
|
for column in entries:
|
|
columns[column].append(row)
|
|
pointers = [0]
|
|
indices = []
|
|
for column in columns:
|
|
indices.extend(column)
|
|
pointers.append(len(indices))
|
|
return [f'const int {name}[{len(values)}] = {{'+','.join(map(str, values))+'};'
|
|
for name, values in [('model_jacobian_column_color', self.colors),
|
|
('model_jacobian_col_ptr', pointers),
|
|
('model_jacobian_row_index', indices)]]
|
|
|
|
|
|
class StateDependencies:
|
|
def __init__(self, state_count):
|
|
self.state_count = state_count
|
|
self.seeds = {f'y[{i}]': 1 << i for i in range(state_count)}
|
|
self.inputs = {}
|
|
|
|
def assign(self, target, inputs):
|
|
# Union repeated writes, rather than dropping dependencies from another
|
|
# branch or an earlier in-place value (stop motion is mode dependent).
|
|
self.inputs.setdefault(target, set()).update(inputs)
|
|
|
|
def expression(self, target, expression):
|
|
self.assign(target, expression_inputs(expression))
|
|
|
|
def project_states(self, offsets):
|
|
refs = {f'y[{i}]' for i in offsets}
|
|
for target in refs:
|
|
self.assign(target, refs)
|
|
|
|
def stop_motion(self, velocity_index, position_index):
|
|
refs = {f'y[{velocity_index}]', f'y[{position_index}]',
|
|
f'dy[{velocity_index}]', f'dy[{position_index}]'}
|
|
for target in (f'dy[{velocity_index}]', f'dy[{position_index}]'):
|
|
self.assign(target, refs)
|
|
|
|
def computation(self, operation):
|
|
# A schedule SCC reaches a fixed point in build(), so every member gains
|
|
# every external state dependency even through pressure/stream loops.
|
|
for output in operation.outputs:
|
|
self.assign(output, operation.inputs)
|
|
|
|
def build(self):
|
|
if self.state_count == 0:
|
|
return JacobianStructure.dense(1, 'Algebraic-only internal state uses default differences')
|
|
required = {f'dy[{i}]' for i in range(self.state_count)}
|
|
pending = list(required)
|
|
reachable = set()
|
|
missing = set()
|
|
while pending:
|
|
key = pending.pop()
|
|
if key in reachable:
|
|
continue
|
|
reachable.add(key)
|
|
if key not in self.inputs and key not in self.seeds:
|
|
missing.add(key)
|
|
pending.extend(self.inputs.get(key, ()))
|
|
if missing:
|
|
return JacobianStructure.dense(self.state_count,
|
|
'Unresolved structural inputs: '+', '.join(sorted(missing)[:16]))
|
|
consumers = {key: set() for key in reachable}
|
|
for target in reachable:
|
|
for source in self.inputs.get(target, ()):
|
|
consumers[source].add(target)
|
|
masks = dict(self.seeds)
|
|
queue = deque(key for key in self.seeds if key in reachable)
|
|
queued = set(queue)
|
|
while queue:
|
|
source = queue.popleft()
|
|
queued.remove(source)
|
|
for target in consumers[source]:
|
|
merged = masks.get(target, 0) | masks[source]
|
|
if merged != masks.get(target, 0):
|
|
masks[target] = merged
|
|
if target not in queued:
|
|
queue.append(target)
|
|
queued.add(target)
|
|
# A diagonal overestimate is safe and keeps isolated/dummy rows valid
|
|
# without introducing zero-length C arrays or uncolored state columns.
|
|
rows = tuple(tuple(column for column in range(self.state_count)
|
|
if (masks.get(f'dy[{row}]', 0) | (1 << row)) >> column & 1)
|
|
for row in range(self.state_count))
|
|
conflicts = [set() for _ in range(self.state_count)]
|
|
for entries in rows:
|
|
for column in entries:
|
|
conflicts[column].update(set(entries) - {column})
|
|
colors = {}
|
|
while len(colors) < self.state_count:
|
|
column = max((i for i in range(self.state_count) if i not in colors),
|
|
key=lambda i: (len({colors[j] for j in conflicts[i] if j in colors}),
|
|
len(conflicts[i]), -i))
|
|
used = {colors[j] for j in conflicts[column] if j in colors}
|
|
color = 0
|
|
while color in used:
|
|
color += 1
|
|
colors[column] = color
|
|
ordered = tuple(colors[i] for i in range(self.state_count))
|
|
for entries in rows:
|
|
if len({ordered[column] for column in entries}) != len(entries):
|
|
raise AssertionError('Jacobian coloring contains a row conflict')
|
|
return JacobianStructure(self.state_count, rows, ordered)
|