Skip to content

Commit 649e68a

Browse files
nachawatinachawati
committed
feat(core): vectorized SIMD-style tape operations for efficient array AD/IA/McCormick evaluation
1 parent 1c807e9 commit 649e68a

7 files changed

Lines changed: 335 additions & 3 deletions

File tree

‎packages/core/src/compiler/modelica/ad-codegen.ts‎

Lines changed: 102 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,19 @@ export type TapeOp =
2828
| { type: "tan"; a: number }
2929
| { type: "exp"; a: number }
3030
| { type: "log"; a: number }
31-
| { type: "sqrt"; a: number };
31+
| { type: "sqrt"; a: number }
32+
// ── Vector ops: SIMD-style operations on contiguous blocks ──
33+
// These occupy `size` consecutive tape slots starting at the op's index.
34+
| { type: "vec_var"; baseName: string; size: number }
35+
| { type: "vec_const"; vals: number[]; size: number }
36+
| { type: "vec_add"; a: number; b: number; size: number }
37+
| { type: "vec_sub"; a: number; b: number; size: number }
38+
| { type: "vec_mul"; a: number; b: number; size: number }
39+
| { type: "vec_neg"; a: number; size: number }
40+
// Extracts scalar t[a + offset] from a vector block starting at tape index `a`
41+
| { type: "vec_subscript"; a: number; offset: number }
42+
// No-op placeholder for vector padding slots (evaluators skip these)
43+
| { type: "nop" };
3244

3345
function formatCDouble(v: number): string {
3446
if (!isFinite(v)) return v === Infinity ? "INFINITY" : v === -Infinity ? "(-INFINITY)" : "NAN";
@@ -187,6 +199,69 @@ export class StaticTapeBuilder {
187199
return [this.walk(expr)];
188200
}
189201

202+
/**
203+
* Push a vector op that reserves `size` consecutive tape slots.
204+
* The op itself occupies slot `idx`, and `size-1` placeholder slots follow.
205+
* Returns the starting index of the vector block.
206+
*/
207+
public pushVecOp(op: TapeOp & { size: number }): number {
208+
const idx = this.ops.length;
209+
this.ops.push(op);
210+
// Reserve remaining size-1 slots with nop placeholders
211+
for (let i = 1; i < op.size; i++) {
212+
this.ops.push({ type: "nop" }); // Placeholder; evaluators skip these
213+
}
214+
return idx;
215+
}
216+
217+
/**
218+
* Walk an array expression as a single vectorized block.
219+
* If `expr` is a `ModelicaArray` of named variables, emits `vec_var` etc.
220+
* Returns the starting tape index and the block size.
221+
*
222+
* For array binary ops like `{a,b} + {c,d}`, emits `vec_add`.
223+
* Falls back to element-wise scalar tapes for complex expressions.
224+
*/
225+
public walkArrayVectorized(expr: ModelicaExpression): { startIdx: number; size: number } {
226+
if (!(expr instanceof ModelicaArray)) {
227+
// Scalar: single tape slot
228+
const idx = this.walk(expr);
229+
return { startIdx: idx, size: 1 };
230+
}
231+
232+
const flatElems = [...expr.flatElements];
233+
const size = flatElems.length;
234+
if (size === 0) return { startIdx: this.pushOp({ type: "const", val: 0 }), size: 1 };
235+
236+
// Check if all elements are simple named variables with a common base
237+
const allNamed = flatElems.every((e) => e instanceof ModelicaNameExpression);
238+
if (allNamed && size > 1) {
239+
// Extract common base name (e.g., "x[1]", "x[2]" → base = "x")
240+
const firstName = (flatElems[0] as ModelicaNameExpression).name;
241+
const bracketPos = firstName.indexOf("[");
242+
const baseName = bracketPos >= 0 ? firstName.substring(0, bracketPos) : firstName;
243+
const startIdx = this.pushVecOp({ type: "vec_var", baseName, size });
244+
return { startIdx, size };
245+
}
246+
247+
// Check if all elements are constants
248+
const allConst = flatElems.every((e) => e instanceof ModelicaRealLiteral || e instanceof ModelicaIntegerLiteral);
249+
if (allConst && size > 1) {
250+
const vals = flatElems.map((e) =>
251+
e instanceof ModelicaRealLiteral ? e.value : e instanceof ModelicaIntegerLiteral ? e.value : 0,
252+
);
253+
const startIdx = this.pushVecOp({ type: "vec_const", vals, size });
254+
return { startIdx, size };
255+
}
256+
257+
// Fallback: walk each element as scalar, pack into contiguous block
258+
const startIdx = this.ops.length;
259+
for (const elem of flatElems) {
260+
this.walk(elem);
261+
}
262+
return { startIdx, size };
263+
}
264+
190265
/**
191266
* Emit the C-code evaluating the expressions step-by-step.
192267
* `varResolver` maps the Modelica variable name into a valid C-code getter string
@@ -264,6 +339,32 @@ export class StaticTapeBuilder {
264339
case "sqrt":
265340
rhs = `sqrt(t[${op.a}])`;
266341
break;
342+
// ── Vector ops: emit for-loops ──
343+
case "vec_var":
344+
lines.push(` for (int _k = 0; _k < ${op.size}; _k++) t[${i}+_k] = ${varResolver(`${op.baseName}[_k+1]`)};`);
345+
continue; // Skip default t[i] = rhs line
346+
case "vec_const":
347+
for (let k = 0; k < op.size; k++) {
348+
lines.push(` t[${i + k}] = ${formatCDouble(op.vals[k] ?? 0)};`);
349+
}
350+
continue;
351+
case "vec_add":
352+
lines.push(` for (int _k = 0; _k < ${op.size}; _k++) t[${i}+_k] = t[${op.a}+_k] + t[${op.b}+_k];`);
353+
continue;
354+
case "vec_sub":
355+
lines.push(` for (int _k = 0; _k < ${op.size}; _k++) t[${i}+_k] = t[${op.a}+_k] - t[${op.b}+_k];`);
356+
continue;
357+
case "vec_mul":
358+
lines.push(` for (int _k = 0; _k < ${op.size}; _k++) t[${i}+_k] = t[${op.a}+_k] * t[${op.b}+_k];`);
359+
continue;
360+
case "vec_neg":
361+
lines.push(` for (int _k = 0; _k < ${op.size}; _k++) t[${i}+_k] = -t[${op.a}+_k];`);
362+
continue;
363+
case "vec_subscript":
364+
rhs = `t[${op.a + op.offset}]`;
365+
break;
366+
case "nop":
367+
continue; // Skip padding slots
267368
}
268369
lines.push(` t[${i}] = ${rhs};`);
269370
}

‎packages/core/src/compiler/modelica/ad-jacobian.ts‎

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,42 @@ export function evaluateTapeForward(ops: TapeOp[], varValues: Map<string, number
7878
case "sqrt":
7979
t[i] = Math.sqrt(t[op.a] ?? 0);
8080
break;
81+
// ── Vector ops ──
82+
case "vec_var":
83+
for (let k = 0; k < op.size; k++) {
84+
t[i + k] = varValues.get(`${op.baseName}[${k + 1}]`) ?? 0;
85+
}
86+
break;
87+
case "vec_const":
88+
for (let k = 0; k < op.size; k++) {
89+
t[i + k] = op.vals[k] ?? 0;
90+
}
91+
break;
92+
case "vec_add":
93+
for (let k = 0; k < op.size; k++) {
94+
t[i + k] = (t[op.a + k] ?? 0) + (t[op.b + k] ?? 0);
95+
}
96+
break;
97+
case "vec_sub":
98+
for (let k = 0; k < op.size; k++) {
99+
t[i + k] = (t[op.a + k] ?? 0) - (t[op.b + k] ?? 0);
100+
}
101+
break;
102+
case "vec_mul":
103+
for (let k = 0; k < op.size; k++) {
104+
t[i + k] = (t[op.a + k] ?? 0) * (t[op.b + k] ?? 0);
105+
}
106+
break;
107+
case "vec_neg":
108+
for (let k = 0; k < op.size; k++) {
109+
t[i + k] = -(t[op.a + k] ?? 0);
110+
}
111+
break;
112+
case "vec_subscript":
113+
t[i] = t[op.a + op.offset] ?? 0;
114+
break;
115+
case "nop":
116+
break;
81117
}
82118
}
83119
return t;
@@ -140,6 +176,38 @@ export function evaluateTapeReverse(ops: TapeOp[], t: Float64Array, outputIndex:
140176
case "sqrt":
141177
dt[op.a] = (dt[op.a] ?? 0) + dti / (2 * (t[i] ?? 1));
142178
break;
179+
// ── Vector ops reverse ──
180+
case "vec_add":
181+
for (let k = 0; k < op.size; k++) {
182+
const dk = dt[i + k] ?? 0;
183+
dt[op.a + k] = (dt[op.a + k] ?? 0) + dk;
184+
dt[op.b + k] = (dt[op.b + k] ?? 0) + dk;
185+
}
186+
break;
187+
case "vec_sub":
188+
for (let k = 0; k < op.size; k++) {
189+
const dk = dt[i + k] ?? 0;
190+
dt[op.a + k] = (dt[op.a + k] ?? 0) + dk;
191+
dt[op.b + k] = (dt[op.b + k] ?? 0) - dk;
192+
}
193+
break;
194+
case "vec_mul":
195+
for (let k = 0; k < op.size; k++) {
196+
const dk = dt[i + k] ?? 0;
197+
dt[op.a + k] = (dt[op.a + k] ?? 0) + dk * (t[op.b + k] ?? 0);
198+
dt[op.b + k] = (dt[op.b + k] ?? 0) + dk * (t[op.a + k] ?? 0);
199+
}
200+
break;
201+
case "vec_neg":
202+
for (let k = 0; k < op.size; k++) {
203+
dt[op.a + k] = (dt[op.a + k] ?? 0) - (dt[i + k] ?? 0);
204+
}
205+
break;
206+
case "vec_subscript":
207+
dt[op.a + op.offset] = (dt[op.a + op.offset] ?? 0) + dti;
208+
break;
209+
case "nop":
210+
break;
143211
}
144212
}
145213

@@ -149,6 +217,11 @@ export function evaluateTapeReverse(ops: TapeOp[], t: Float64Array, outputIndex:
149217
const op = ops[i]!; // eslint-disable-line @typescript-eslint/no-non-null-assertion
150218
if (op.type === "var") {
151219
gradients.set(op.name, (gradients.get(op.name) ?? 0) + (dt[i] ?? 0));
220+
} else if (op.type === "vec_var") {
221+
for (let k = 0; k < op.size; k++) {
222+
const name = `${op.baseName}[${k + 1}]`;
223+
gradients.set(name, (gradients.get(name) ?? 0) + (dt[i + k] ?? 0));
224+
}
152225
}
153226
}
154227
return gradients;

‎packages/core/src/compiler/modelica/branch-and-bound.ts‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -395,3 +395,28 @@ export function buildSbbFromDAE(
395395
constraintTapes,
396396
};
397397
}
398+
399+
/**
400+
* Expand array variable bounds in a DomainBox.
401+
* Given a box with `x → [lo, hi]` and a variable `x` with arrayDimensions = [3],
402+
* produces `x[1] → [lo,hi]`, `x[2] → [lo,hi]`, `x[3] → [lo,hi]`.
403+
*
404+
* @param box Initial domain box (may contain array-valued variable names)
405+
* @param dae DAE for variable metadata (arrayDimensions)
406+
* @returns Expanded domain box with per-element entries
407+
*/
408+
export function expandArrayBounds(box: DomainBox, dae: ModelicaDAE): DomainBox {
409+
const expanded: DomainBox = new Map();
410+
for (const [name, interval] of box) {
411+
const v = dae.variables.find((dv) => dv.name === name);
412+
if (v?.arrayDimensions && v.arrayDimensions.length > 0) {
413+
const size = v.arrayDimensions.reduce((a: number, b: number) => a * b, 1);
414+
for (let i = 0; i < size; i++) {
415+
expanded.set(`${name}[${i + 1}]`, new Interval(interval.lo, interval.hi));
416+
}
417+
} else {
418+
expanded.set(name, interval);
419+
}
420+
}
421+
return expanded;
422+
}

‎packages/core/src/compiler/modelica/fmu-codegen.ts‎

Lines changed: 23 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,8 @@
1717
import { StaticTapeBuilder } from "./ad-codegen.js";
1818
import type { ModelicaDAE, ModelicaExpression } from "./dae.js";
1919
import {
20+
ModelicaArray,
21+
ModelicaArrayEquation,
2022
ModelicaBinaryExpression,
2123
ModelicaBooleanLiteral,
2224
ModelicaFunctionCallEquation,
@@ -618,12 +620,24 @@ function generateInitializeSolve(id: string, dae: ModelicaDAE, result: FmuResult
618620
return lines;
619621
}
620622

621-
// Collect initial equations that are simple equations (LHS = RHS)
623+
// Collect initial equations — unroll array equations into per-element scalar equations
622624
const initEqs: { lhs: ModelicaExpression; rhs: ModelicaExpression }[] = [];
623625
for (const eq of dae.initialEquations) {
624626
if ("expression1" in eq && "expression2" in eq) {
625627
const se = eq as { expression1: ModelicaExpression; expression2: ModelicaExpression };
626-
initEqs.push({ lhs: se.expression1, rhs: se.expression2 });
628+
if (eq instanceof ModelicaArrayEquation) {
629+
// Unroll array equation into per-element scalar equations
630+
const lhsElems = se.expression1 instanceof ModelicaArray ? [...se.expression1.flatElements] : [se.expression1];
631+
const rhsElems = se.expression2 instanceof ModelicaArray ? [...se.expression2.flatElements] : [se.expression2];
632+
const n = Math.max(lhsElems.length, rhsElems.length);
633+
for (let i = 0; i < n; i++) {
634+
const lhs = lhsElems[i] ?? lhsElems[0];
635+
const rhs = rhsElems[i] ?? rhsElems[0];
636+
if (lhs && rhs) initEqs.push({ lhs, rhs });
637+
}
638+
} else {
639+
initEqs.push({ lhs: se.expression1, rhs: se.expression2 });
640+
}
627641
}
628642
}
629643

@@ -2193,6 +2207,13 @@ function extractDerName(expr: unknown): string | null {
21932207
/** Recursively collect all variable names referenced in an expression. */
21942208
function collectReferencedNames(expr: unknown, names: Set<string>): void {
21952209
if (!expr || typeof expr !== "object") return;
2210+
// Handle ModelicaArray — recurse into each element
2211+
if (expr instanceof ModelicaArray) {
2212+
for (const elem of expr.flatElements) {
2213+
collectReferencedNames(elem, names);
2214+
}
2215+
return;
2216+
}
21962217
if ("name" in expr) {
21972218
const nameVal = (expr as { name: unknown }).name;
21982219
if (typeof nameVal === "string") {

‎packages/core/src/compiler/modelica/init-solver.ts‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,8 @@
1212
import { StaticTapeBuilder, type TapeOp } from "./ad-codegen.js";
1313
import { evaluateTapeForward, evaluateTapeReverse } from "./ad-jacobian.js";
1414
import {
15+
ModelicaArray,
16+
ModelicaArrayEquation,
1517
ModelicaBooleanLiteral,
1618
type ModelicaDAE,
1719
type ModelicaEquation,
@@ -58,6 +60,13 @@ function collectVarNames(expr: ModelicaExpression, names: Set<string>): void {
5860
names.add(expr.name);
5961
return;
6062
}
63+
// Recurse into array elements
64+
if (expr instanceof ModelicaArray) {
65+
for (const elem of expr.flatElements) {
66+
collectVarNames(elem, names);
67+
}
68+
return;
69+
}
6170
if ("name" in expr) names.add((expr as { name: string }).name);
6271
if ("operand" in expr) collectVarNames((expr as { operand: ModelicaExpression }).operand, names);
6372
if ("operand1" in expr) collectVarNames((expr as { operand1: ModelicaExpression }).operand1, names);
@@ -283,10 +292,38 @@ export function solveInitialEquations(
283292
// Classify equations
284293
const classified: ClassifiedEq[] = [];
285294
for (const eq of dae.initialEquations) {
295+
if (eq instanceof ModelicaArrayEquation) {
296+
// Unroll array initial equation into per-element implicit equations
297+
const se = eq as { expression1: ModelicaExpression; expression2: ModelicaExpression };
298+
const lhsElems = se.expression1 instanceof ModelicaArray ? [...se.expression1.flatElements] : [se.expression1];
299+
const rhsElems = se.expression2 instanceof ModelicaArray ? [...se.expression2.flatElements] : [se.expression2];
300+
const n = Math.max(lhsElems.length, rhsElems.length);
301+
for (let i = 0; i < n; i++) {
302+
const lhs = lhsElems[i] ?? lhsElems[0];
303+
const rhs = rhsElems[i] ?? rhsElems[0];
304+
if (lhs && rhs) {
305+
// Create a synthetic simple equation for each element
306+
const elemEq = { expression1: lhs, expression2: rhs } as unknown as ModelicaEquation;
307+
const c = classifyInitialEquation(elemEq, unknowns);
308+
if (c) classified.push(c);
309+
}
310+
}
311+
continue;
312+
}
286313
const c = classifyInitialEquation(eq, unknowns);
287314
if (c) classified.push(c);
288315
}
289316

317+
// Expand array variable unknowns: for vars with arrayDimensions, add subscripted names
318+
for (const v of dae.variables) {
319+
if (v.arrayDimensions && v.arrayDimensions.length > 0 && unknowns.has(v.name)) {
320+
const size = v.arrayDimensions.reduce((a: number, b: number) => a * b, 1);
321+
for (let i = 0; i < size; i++) {
322+
unknowns.add(`${v.name}[${i + 1}]`);
323+
}
324+
}
325+
}
326+
290327
// Phase 1: Solve explicit equations first (direct assignment)
291328
const explicitEqs = classified.filter((c): c is ExplicitEq => c.type === "explicit");
292329
const implicitEqs = classified.filter((c): c is ImplicitEq => c.type === "implicit");

0 commit comments

Comments
 (0)