gradient-script 0.2.0 → 0.3.1

package/dist/dsl/egraph/Optimizer.d.ts

@@ -0,0 +1,50 @@
+/**
+ * E-Graph Optimizer for GradientScript
+ *
+ * Main entry point for e-graph-based optimization.
+ * Uses equality saturation for CSE and algebraic simplification.
+ */
+import { Expression } from '../AST.js';
+import { SaturationStats } from './Rewriter.js';
+import { CostModel } from './Extractor.js';
+/**
+ * Result from e-graph optimization (matches GlobalCSEResult interface)
+ */
+export interface EGraphOptimizeResult {
+    intermediates: Map<string, Expression>;
+    gradients: Map<string, Map<string, Expression>>;
+    stats?: OptimizationStats;
+}
+/**
+ * Statistics from optimization
+ */
+export interface OptimizationStats {
+    saturation: SaturationStats;
+    tempsCreated: number;
+    totalCost: number;
+}
+/**
+ * Options for e-graph optimization
+ */
+export interface EGraphOptimizeOptions {
+    /** Maximum saturation iterations (default: 30) */
+    maxIterations?: number;
+    /** Maximum e-graph size before stopping saturation (default: 15000) */
+    maxClassSize?: number;
+    /** Which rule sets to use (default: ['core', 'algebra']) */
+    ruleSets?: ('core' | 'algebra' | 'function')[];
+    /** Use phased saturation (core first, then others) */
+    phased?: boolean;
+    /** Minimum cost for a subexpression to become a temp (default: 3) */
+    minSharedCost?: number;
+    /** Custom cost model */
+    costModel?: CostModel;
+    /** Print verbose output */
+    verbose?: boolean;
+}
+/**
+ * Optimize gradients using e-graph equality saturation
+ *
+ * This is designed to be a drop-in replacement for eliminateCommonSubexpressionsGlobal
+ */
+export declare function optimizeWithEGraph(allGradients: Map<string, Map<string, Expression>>, options?: EGraphOptimizeOptions): EGraphOptimizeResult;

package/dist/dsl/egraph/Optimizer.js

@@ -0,0 +1,88 @@
+/**
+ * E-Graph Optimizer for GradientScript
+ *
+ * Main entry point for e-graph-based optimization.
+ * Uses equality saturation for CSE and algebraic simplification.
+ */
+import { EGraph } from './EGraph.js';
+import { addGradients, getRootIds } from './Convert.js';
+import { saturate, saturatePhased } from './Rewriter.js';
+import { extractWithCSE, defaultCostModel } from './Extractor.js';
+import { coreRules, algebraRules, functionRules } from './Rules.js';
+/**
+ * Optimize gradients using e-graph equality saturation
+ *
+ * This is designed to be a drop-in replacement for eliminateCommonSubexpressionsGlobal
+ */
+export function optimizeWithEGraph(allGradients, options = {}) {
+    const { maxIterations = 30, maxClassSize = 15000, ruleSets = ['core', 'algebra'], phased = true, minSharedCost = 2, costModel = defaultCostModel, verbose = false } = options;
+    // Build e-graph from all gradients
+    const egraph = new EGraph();
+    const gradientIds = addGradients(egraph, allGradients);
+    const rootIds = getRootIds(gradientIds);
+    if (verbose) {
+        console.log(`[egraph] Added ${egraph.size} e-classes from ${rootIds.length} gradient expressions`);
+    }
+    // Select rules
+    const rules = selectRules(ruleSets);
+    // Saturate
+    let stats;
+    if (phased) {
+        const phases = [];
+        if (ruleSets.includes('core'))
+            phases.push(coreRules);
+        if (ruleSets.includes('algebra'))
+            phases.push(algebraRules);
+        if (ruleSets.includes('function'))
+            phases.push(functionRules);
+        stats = saturatePhased(egraph, phases, { maxIterations, maxClassSize, verbose });
+    }
+    else {
+        stats = saturate(egraph, rules, { maxIterations, maxClassSize, verbose });
+    }
+    if (verbose) {
+        console.log(`[egraph] Saturation: ${stats.iterations} iters, ${stats.merges} merges, ${egraph.size} classes`);
+    }
+    // Extract with CSE
+    const extraction = extractWithCSE(egraph, rootIds, costModel, minSharedCost);
+    if (verbose) {
+        console.log(`[egraph] Extracted ${extraction.temps.size} temps, total cost ${extraction.totalCost}`);
+    }
+    // Convert back to gradient structure
+    const optimizedGradients = new Map();
+    let rootIndex = 0;
+    for (const [paramName, componentIds] of gradientIds) {
+        const components = new Map();
+        for (const [comp] of componentIds) {
+            const rootId = rootIds[rootIndex++];
+            const expr = extraction.expressions.get(rootId);
+            if (!expr) {
+                throw new Error(`Missing extraction for root ${rootId}`);
+            }
+            components.set(comp, expr);
+        }
+        optimizedGradients.set(paramName, components);
+    }
+    return {
+        intermediates: extraction.temps,
+        gradients: optimizedGradients,
+        stats: {
+            saturation: stats,
+            tempsCreated: extraction.temps.size,
+            totalCost: extraction.totalCost
+        }
+    };
+}
+/**
+ * Select rules based on rule set names
+ */
+function selectRules(ruleSets) {
+    const rules = [];
+    if (ruleSets.includes('core'))
+        rules.push(...coreRules);
+    if (ruleSets.includes('algebra'))
+        rules.push(...algebraRules);
+    if (ruleSets.includes('function'))
+        rules.push(...functionRules);
+    return rules;
+}

package/dist/dsl/egraph/Pattern.d.ts

@@ -0,0 +1,80 @@
+/**
+ * Pattern Matching for E-Graph Rewrite Rules
+ *
+ * Patterns are expression templates with variables (?a, ?b, etc.)
+ * that can match against e-classes in the e-graph.
+ */
+import { EGraph } from './EGraph.js';
+import { EClassId } from './ENode.js';
+/**
+ * Pattern AST
+ */
+export type Pattern = {
+    tag: 'pvar';
+    name: string;
+} | {
+    tag: 'pnum';
+    value: number;
+} | {
+    tag: 'padd';
+    left: Pattern;
+    right: Pattern;
+} | {
+    tag: 'pmul';
+    left: Pattern;
+    right: Pattern;
+} | {
+    tag: 'psub';
+    left: Pattern;
+    right: Pattern;
+} | {
+    tag: 'pdiv';
+    left: Pattern;
+    right: Pattern;
+} | {
+    tag: 'ppow';
+    left: Pattern;
+    right: Pattern;
+} | {
+    tag: 'pneg';
+    child: Pattern;
+} | {
+    tag: 'pinv';
+    child: Pattern;
+} | {
+    tag: 'pcall';
+    name: string;
+    args: Pattern[];
+};
+/**
+ * A substitution mapping pattern variables to e-class IDs
+ */
+export type Substitution = Map<string, EClassId>;
+/**
+ * Parse a pattern string into a Pattern AST
+ *
+ * Syntax:
+ *   ?a, ?b, ?x     - pattern variables
+ *   0, 1, 2, -1    - number literals
+ *   (+ ?a ?b)      - addition
+ *   (* ?a ?b)      - multiplication
+ *   (- ?a ?b)      - subtraction
+ *   (/ ?a ?b)      - division
+ *   (^ ?a ?b)      - power
+ *   (neg ?a)       - negation
+ *   (sqrt ?a)      - function call
+ */
+export declare function parsePattern(input: string): Pattern;
+/**
+ * Match a pattern against an e-class, returning all valid substitutions
+ */
+export declare function matchPattern(egraph: EGraph, pattern: Pattern, classId: EClassId): Substitution[];
+/**
+ * Instantiate a pattern with a substitution, adding nodes to the e-graph
+ * Returns the e-class ID of the instantiated pattern
+ */
+export declare function instantiatePattern(egraph: EGraph, pattern: Pattern, subst: Substitution): EClassId;
+/**
+ * Convert pattern to string for debugging
+ */
+export declare function patternToString(pattern: Pattern): string;

package/dist/dsl/egraph/Pattern.js

@@ -0,0 +1,325 @@
+/**
+ * Pattern Matching for E-Graph Rewrite Rules
+ *
+ * Patterns are expression templates with variables (?a, ?b, etc.)
+ * that can match against e-classes in the e-graph.
+ */
+/**
+ * Parse a pattern string into a Pattern AST
+ *
+ * Syntax:
+ *   ?a, ?b, ?x     - pattern variables
+ *   0, 1, 2, -1    - number literals
+ *   (+ ?a ?b)      - addition
+ *   (* ?a ?b)      - multiplication
+ *   (- ?a ?b)      - subtraction
+ *   (/ ?a ?b)      - division
+ *   (^ ?a ?b)      - power
+ *   (neg ?a)       - negation
+ *   (sqrt ?a)      - function call
+ */
+export function parsePattern(input) {
+    const tokens = tokenize(input);
+    let pos = 0;
+    function peek() {
+        return tokens[pos];
+    }
+    function consume() {
+        return tokens[pos++];
+    }
+    function parseExpr() {
+        const token = peek();
+        if (token === '(') {
+            consume(); // (
+            const op = consume();
+            let result;
+            if (op === 'neg') {
+                const child = parseExpr();
+                result = { tag: 'pneg', child };
+            }
+            else if (op === 'inv') {
+                const child = parseExpr();
+                result = { tag: 'pinv', child };
+            }
+            else if (['+', '-', '*', '/', '^'].includes(op)) {
+                const left = parseExpr();
+                const right = parseExpr();
+                switch (op) {
+                    case '+':
+                        result = { tag: 'padd', left, right };
+                        break;
+                    case '-':
+                        result = { tag: 'psub', left, right };
+                        break;
+                    case '*':
+                        result = { tag: 'pmul', left, right };
+                        break;
+                    case '/':
+                        result = { tag: 'pdiv', left, right };
+                        break;
+                    case '^':
+                        result = { tag: 'ppow', left, right };
+                        break;
+                    default: throw new Error(`Unknown operator: ${op}`);
+                }
+            }
+            else {
+                // Function call
+                const args = [];
+                while (peek() !== ')') {
+                    args.push(parseExpr());
+                }
+                result = { tag: 'pcall', name: op, args };
+            }
+            if (consume() !== ')') {
+                throw new Error('Expected )');
+            }
+            return result;
+        }
+        if (token?.startsWith('?')) {
+            consume();
+            return { tag: 'pvar', name: token.slice(1) };
+        }
+        if (token && /^-?\d+(\.\d+)?$/.test(token)) {
+            consume();
+            return { tag: 'pnum', value: parseFloat(token) };
+        }
+        throw new Error(`Unexpected token: ${token}`);
+    }
+    const result = parseExpr();
+    if (pos < tokens.length) {
+        throw new Error(`Unexpected token after pattern: ${tokens[pos]}`);
+    }
+    return result;
+}
+/**
+ * Tokenize a pattern string
+ */
+function tokenize(input) {
+    const tokens = [];
+    let i = 0;
+    while (i < input.length) {
+        const ch = input[i];
+        if (/\s/.test(ch)) {
+            i++;
+            continue;
+        }
+        if (ch === '(' || ch === ')') {
+            tokens.push(ch);
+            i++;
+            continue;
+        }
+        // Variable or operator or number
+        let token = '';
+        while (i < input.length && !/[\s()]/.test(input[i])) {
+            token += input[i];
+            i++;
+        }
+        if (token) {
+            tokens.push(token);
+        }
+    }
+    return tokens;
+}
+/** Maximum substitutions to return from a single match (prevents explosion) */
+const MAX_SUBSTITUTIONS = 100;
+/**
+ * Match a pattern against an e-class, returning all valid substitutions
+ */
+export function matchPattern(egraph, pattern, classId) {
+    const results = matchPatternWithSubst(egraph, pattern, classId, new Map(), 0);
+    return results.slice(0, MAX_SUBSTITUTIONS);
+}
+const MAX_MATCH_DEPTH = 50;
+function matchPatternWithSubst(egraph, pattern, classId, subst, depth) {
+    if (depth > MAX_MATCH_DEPTH) {
+        return []; // Prevent infinite recursion
+    }
+    const canonId = egraph.find(classId);
+    // Pattern variable - bind or check existing binding
+    if (pattern.tag === 'pvar') {
+        const existing = subst.get(pattern.name);
+        if (existing !== undefined) {
+            // Check if it matches the same e-class
+            if (egraph.find(existing) === canonId) {
+                return [new Map(subst)];
+            }
+            return [];
+        }
+        // Bind the variable
+        const newSubst = new Map(subst);
+        newSubst.set(pattern.name, canonId);
+        return [newSubst];
+    }
+    // Try to match against all nodes in the e-class
+    const nodes = egraph.getNodes(canonId);
+    const results = [];
+    for (const node of nodes) {
+        if (results.length >= MAX_SUBSTITUTIONS)
+            break;
+        const matches = matchNodeWithPattern(egraph, pattern, node, subst, depth);
+        for (const m of matches) {
+            results.push(m);
+            if (results.length >= MAX_SUBSTITUTIONS)
+                break;
+        }
+    }
+    return results;
+}
+function matchNodeWithPattern(egraph, pattern, node, subst, depth) {
+    switch (pattern.tag) {
+        case 'pvar':
+            // Already handled above
+            throw new Error('pvar should be handled in matchPatternWithSubst');
+        case 'pnum':
+            if (node.tag === 'num' && node.value === pattern.value) {
+                return [new Map(subst)];
+            }
+            return [];
+        case 'padd':
+            if (node.tag === 'add') {
+                return matchBinaryChildren(egraph, pattern.left, pattern.right, node.children, subst, depth);
+            }
+            return [];
+        case 'pmul':
+            if (node.tag === 'mul') {
+                return matchBinaryChildren(egraph, pattern.left, pattern.right, node.children, subst, depth);
+            }
+            return [];
+        case 'psub':
+            if (node.tag === 'sub') {
+                return matchBinaryChildren(egraph, pattern.left, pattern.right, node.children, subst, depth);
+            }
+            return [];
+        case 'pdiv':
+            if (node.tag === 'div') {
+                return matchBinaryChildren(egraph, pattern.left, pattern.right, node.children, subst, depth);
+            }
+            return [];
+        case 'ppow':
+            if (node.tag === 'pow') {
+                return matchBinaryChildren(egraph, pattern.left, pattern.right, node.children, subst, depth);
+            }
+            return [];
+        case 'pneg':
+            if (node.tag === 'neg') {
+                return matchPatternWithSubst(egraph, pattern.child, node.child, subst, depth + 1);
+            }
+            return [];
+        case 'pinv':
+            if (node.tag === 'inv') {
+                return matchPatternWithSubst(egraph, pattern.child, node.child, subst, depth + 1);
+            }
+            return [];
+        case 'pcall':
+            if (node.tag === 'call' && node.name === pattern.name && node.children.length === pattern.args.length) {
+                return matchCallChildren(egraph, pattern.args, node.children, subst, depth);
+            }
+            return [];
+    }
+}
+function matchBinaryChildren(egraph, leftPattern, rightPattern, children, subst, depth) {
+    const results = [];
+    // Match left, then right
+    const leftMatches = matchPatternWithSubst(egraph, leftPattern, children[0], subst, depth + 1);
+    for (const leftSubst of leftMatches) {
+        if (results.length >= MAX_SUBSTITUTIONS)
+            break;
+        const rightMatches = matchPatternWithSubst(egraph, rightPattern, children[1], leftSubst, depth + 1);
+        for (const m of rightMatches) {
+            results.push(m);
+            if (results.length >= MAX_SUBSTITUTIONS)
+                break;
+        }
+    }
+    return results;
+}
+function matchCallChildren(egraph, patterns, children, subst, depth) {
+    if (patterns.length === 0) {
+        return [new Map(subst)];
+    }
+    const results = [];
+    const firstMatches = matchPatternWithSubst(egraph, patterns[0], children[0], subst, depth + 1);
+    for (const firstSubst of firstMatches) {
+        if (results.length >= MAX_SUBSTITUTIONS)
+            break;
+        const restMatches = matchCallChildren(egraph, patterns.slice(1), children.slice(1), firstSubst, depth + 1);
+        for (const m of restMatches) {
+            results.push(m);
+            if (results.length >= MAX_SUBSTITUTIONS)
+                break;
+        }
+    }
+    return results;
+}
+/**
+ * Instantiate a pattern with a substitution, adding nodes to the e-graph
+ * Returns the e-class ID of the instantiated pattern
+ */
+export function instantiatePattern(egraph, pattern, subst) {
+    switch (pattern.tag) {
+        case 'pvar': {
+            const id = subst.get(pattern.name);
+            if (id === undefined) {
+                throw new Error(`Unbound pattern variable: ?${pattern.name}`);
+            }
+            return id;
+        }
+        case 'pnum':
+            return egraph.add({ tag: 'num', value: pattern.value });
+        case 'padd': {
+            const left = instantiatePattern(egraph, pattern.left, subst);
+            const right = instantiatePattern(egraph, pattern.right, subst);
+            return egraph.add({ tag: 'add', children: [left, right] });
+        }
+        case 'pmul': {
+            const left = instantiatePattern(egraph, pattern.left, subst);
+            const right = instantiatePattern(egraph, pattern.right, subst);
+            return egraph.add({ tag: 'mul', children: [left, right] });
+        }
+        case 'psub': {
+            const left = instantiatePattern(egraph, pattern.left, subst);
+            const right = instantiatePattern(egraph, pattern.right, subst);
+            return egraph.add({ tag: 'sub', children: [left, right] });
+        }
+        case 'pdiv': {
+            const left = instantiatePattern(egraph, pattern.left, subst);
+            const right = instantiatePattern(egraph, pattern.right, subst);
+            return egraph.add({ tag: 'div', children: [left, right] });
+        }
+        case 'ppow': {
+            const left = instantiatePattern(egraph, pattern.left, subst);
+            const right = instantiatePattern(egraph, pattern.right, subst);
+            return egraph.add({ tag: 'pow', children: [left, right] });
+        }
+        case 'pneg': {
+            const child = instantiatePattern(egraph, pattern.child, subst);
+            return egraph.add({ tag: 'neg', child });
+        }
+        case 'pinv': {
+            const child = instantiatePattern(egraph, pattern.child, subst);
+            return egraph.add({ tag: 'inv', child });
+        }
+        case 'pcall': {
+            const args = pattern.args.map(arg => instantiatePattern(egraph, arg, subst));
+            return egraph.add({ tag: 'call', name: pattern.name, children: args });
+        }
+    }
+}
+/**
+ * Convert pattern to string for debugging
+ */
+export function patternToString(pattern) {
+    switch (pattern.tag) {
+        case 'pvar': return `?${pattern.name}`;
+        case 'pnum': return `${pattern.value}`;
+        case 'padd': return `(+ ${patternToString(pattern.left)} ${patternToString(pattern.right)})`;
+        case 'pmul': return `(* ${patternToString(pattern.left)} ${patternToString(pattern.right)})`;
+        case 'psub': return `(- ${patternToString(pattern.left)} ${patternToString(pattern.right)})`;
+        case 'pdiv': return `(/ ${patternToString(pattern.left)} ${patternToString(pattern.right)})`;
+        case 'ppow': return `(^ ${patternToString(pattern.left)} ${patternToString(pattern.right)})`;
+        case 'pneg': return `(neg ${patternToString(pattern.child)})`;
+        case 'pinv': return `(inv ${patternToString(pattern.child)})`;
+        case 'pcall': return `(${pattern.name} ${pattern.args.map(patternToString).join(' ')})`;
+    }
+}

package/dist/dsl/egraph/Rewriter.d.ts

@@ -0,0 +1,44 @@
+/**
+ * Rewrite Engine for E-Graph Equality Saturation
+ *
+ * Applies rewrite rules until saturation (no new merges) or iteration limit.
+ */
+import { EGraph } from './EGraph.js';
+import { Rule } from './Rules.js';
+/**
+ * Statistics from a saturation run
+ */
+export interface SaturationStats {
+    iterations: number;
+    totalMatches: number;
+    merges: number;
+    saturated: boolean;
+    classCount: number;
+}
+/**
+ * Options for saturation
+ */
+export interface SaturationOptions {
+    maxIterations?: number;
+    maxClassSize?: number;
+    maxMatchesPerIter?: number;
+    verbose?: boolean;
+}
+/**
+ * Apply equality saturation to an e-graph
+ *
+ * Repeatedly applies rewrite rules until:
+ * - No new equivalences are discovered (saturated)
+ * - Max iterations reached
+ * - E-graph size limit exceeded
+ */
+export declare function saturate(egraph: EGraph, rules: Rule[], options?: SaturationOptions): SaturationStats;
+/**
+ * Apply a single rule once, returning number of merges
+ */
+export declare function applyRuleOnce(egraph: EGraph, rule: Rule): number;
+/**
+ * Apply rules in phases for better control
+ * E.g., apply core rules first, then algebra rules
+ */
+export declare function saturatePhased(egraph: EGraph, phases: Rule[][], options?: SaturationOptions): SaturationStats;