feat(core): vectorize SIMD-style tape operations for AD, IA, and McCormick Array evaluate speedup

nachawati · nachawati · commit 980eabcabbbe · 2026-03-30T10:14:03.000+03:00
diff --git a/packages/core/src/compiler/modelica/ad-codegen.ts b/packages/core/src/compiler/modelica/ad-codegen.ts
@@ -534,7 +534,14 @@ export class StaticTapeBuilder {
     // Reverse topological traversal
     for (let i = this.ops.length - 1; i >= 0; i--) {
       const op = this.ops[i]!; // eslint-disable-line @typescript-eslint/no-non-null-assertion
-      if (op.type === "const" || op.type === "var") continue;
+      if (
+        op.type === "const" ||
+        op.type === "var" ||
+        op.type === "vec_var" ||
+        op.type === "vec_const" ||
+        op.type === "nop"
+      )
+        continue;
 
       // Optimization: if dt is 0, skip branching
       lines.push(`  if (dt[${i}] != 0.0) {`);
@@ -580,6 +587,29 @@ export class StaticTapeBuilder {
         case "sqrt":
           lines.push(`    dt[${op.a}] += dt[${i}] / (2.0 * t[${i}]);`);
           break;
+        // ── Vector ops ──
+        case "vec_add":
+          lines.push(`    for (int _k = 0; _k < ${op.size}; _k++) {`);
+          lines.push(`      dt[${op.a}+_k] += dt[${i}+_k]; dt[${op.b}+_k] += dt[${i}+_k];`);
+          lines.push(`    }`);
+          break;
+        case "vec_sub":
+          lines.push(`    for (int _k = 0; _k < ${op.size}; _k++) {`);
+          lines.push(`      dt[${op.a}+_k] += dt[${i}+_k]; dt[${op.b}+_k] -= dt[${i}+_k];`);
+          lines.push(`    }`);
+          break;
+        case "vec_mul":
+          lines.push(`    for (int _k = 0; _k < ${op.size}; _k++) {`);
+          lines.push(`      dt[${op.a}+_k] += dt[${i}+_k] * t[${op.b}+_k];`);
+          lines.push(`      dt[${op.b}+_k] += dt[${i}+_k] * t[${op.a}+_k];`);
+          lines.push(`    }`);
+          break;
+        case "vec_neg":
+          lines.push(`    for (int _k = 0; _k < ${op.size}; _k++) dt[${op.a}+_k] -= dt[${i}+_k];`);
+          break;
+        case "vec_subscript":
+          lines.push(`    dt[${op.a + op.offset}] += dt[${i}];`);
+          break;
       }
       lines.push(`  }`);
     }
@@ -592,6 +622,10 @@ export class StaticTapeBuilder {
         // Variables can be read multiple times (e.g. x in x + x).
         // Since `cache` deduplicates, there is only one `var: x` operator!
         gradients.set(op.name, i);
+      } else if (op.type === "vec_var") {
+        for (let k = 0; k < op.size; k++) {
+          gradients.set(`${op.baseName}[${k + 1}]`, i + k);
+        }
       }
     }
 
diff --git a/packages/core/src/compiler/modelica/interval.ts b/packages/core/src/compiler/modelica/interval.ts
@@ -416,6 +416,45 @@ export function emitIntervalForwardC(
       case "sqrt":
         lines.push(`t_lo[${i}] = sqrt(fmax(0.0,t_lo[${op.a}])); t_hi[${i}] = sqrt(fmax(0.0,t_hi[${op.a}]));`);
         break;
+      // ── Vector ops ──
+      case "vec_var":
+        for (let k = 0; k < op.size; k++) {
+          const vr = varResolver(`${op.baseName}[${k + 1}]`);
+          lines.push(`t_lo[${i + k}] = ${vr.lo}; t_hi[${i + k}] = ${vr.hi};`);
+        }
+        break;
+      case "vec_const":
+        for (let k = 0; k < op.size; k++) {
+          lines.push(`t_lo[${i + k}] = ${formatNum(op.vals[k] ?? 0)}; t_hi[${i + k}] = ${formatNum(op.vals[k] ?? 0)};`);
+        }
+        break;
+      case "vec_add":
+        lines.push(
+          `for (int _k = 0; _k < ${op.size}; _k++) { t_lo[${i}+_k] = t_lo[${op.a}+_k] + t_lo[${op.b}+_k]; t_hi[${i}+_k] = t_hi[${op.a}+_k] + t_hi[${op.b}+_k]; }`,
+        );
+        break;
+      case "vec_sub":
+        lines.push(
+          `for (int _k = 0; _k < ${op.size}; _k++) { t_lo[${i}+_k] = t_lo[${op.a}+_k] - t_hi[${op.b}+_k]; t_hi[${i}+_k] = t_hi[${op.a}+_k] - t_lo[${op.b}+_k]; }`,
+        );
+        break;
+      case "vec_mul":
+        lines.push(`for (int _k = 0; _k < ${op.size}; _k++) {`);
+        lines.push(`  double p1=t_lo[${op.a}+_k]*t_lo[${op.b}+_k], p2=t_lo[${op.a}+_k]*t_hi[${op.b}+_k],`);
+        lines.push(`         p3=t_hi[${op.a}+_k]*t_lo[${op.b}+_k], p4=t_hi[${op.a}+_k]*t_hi[${op.b}+_k];`);
+        lines.push(`  t_lo[${i}+_k] = fmin(fmin(p1,p2),fmin(p3,p4)); t_hi[${i}+_k] = fmax(fmax(p1,p2),fmax(p3,p4));`);
+        lines.push(`}`);
+        break;
+      case "vec_neg":
+        lines.push(
+          `for (int _k = 0; _k < ${op.size}; _k++) { t_lo[${i}+_k] = -t_hi[${op.a}+_k]; t_hi[${i}+_k] = -t_lo[${op.a}+_k]; }`,
+        );
+        break;
+      case "vec_subscript":
+        lines.push(`t_lo[${i}] = t_lo[${op.a + op.offset}]; t_hi[${i}] = t_hi[${op.a + op.offset}];`);
+        break;
+      case "nop":
+        break;
     }
   }
   return lines;
