feat(core): vectorized SIMD-style tape operations for efficient array… · modelscript/modelscript@649e68a · GitHub
Skip to content

Commit 649e68a

Browse files
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

‎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

0 commit comments

Comments
 (0)