Hook into Symbolics

[DhairyaLGandhi]

Now that we have some passes written out, it will make sense to hook these into Symbolics (and then MTK/ System) so we can get some feedback on the codegen to improve from there.

One of the ideas is for SymbolicCompilerPasses to "declare" all the rules it has, and then Symbolics can define levels that it wants to target.

Something like the following:

const OPTIMIZATION_LEVELS = Dict(
      :none => Symbol[],
      :basic => [MATMUL_ADD_RULE, LDIV_RULE],
      :standard => [MATMUL_ADD_RULE, LDIV_RULE, TRIU_RULE, TRIL_RULE, ],
      :aggressive => [MATMUL_ADD_RULE, LDIV_RULE, TRIU_RULE, TRIL_RULE, HVNCAT_STATIC_RULE, ORTHO_INV_RULE],
  )

And build_function can then look like

build_function(..., optimize = :none) # default for backwards compat
build_function(..., optimize = :standard)
... 

With hooks such that we can pass:

build_function(..., optimize = [MATMUL_ADD_RULE, ORTHO_INV_RULE])

This would be needed when we want to be careful about what kinds of optimizations we can run on the code, plus in the future it could allow for "config" options in the rules (think choosing factorization method for LDIV).

Lastly this would take in the IR produced post CSE and run through

function optimize_ir(expr::Code.Let, state::Code.CSEState, rules)
    for rule in sort(rules, by = x -> x.priority)
        optimized = apply_optimization_rules(expr, state, rule)
        if optimized !== nothing
            expr = optimized
        end
    end
    return expr
end

cc @AayushSabharwal

[DhairyaLGandhi]

Similarly

const BASIC_PASSES = [                                                                                                                                                            
      MATMUL_ADD_RULE,                                                                                                                                                              
      LDIV_RULE,                                                                                                                                                                    
  ]                                                                                                                                                                                 
                                                                                                                                                                                    
  const STANDARD_PASSES = [                                                                                                                                                         
      BASIC_PASSES;                                                                                                                                                                    
      INPLACE_RULE;                                                                                                                                                            
      TRIU_RULE;
      TRIL_RULE                                                                                                                                                    
  ]                                                                                                                                                                                 
                                                                                                                                                                                    
  const AGGRESSIVE_PASSES = [                                                                                                                                                       
      STANDARD_PASSES;
      HVNCAT_STATIC_RULE;
      ORTHO_INV_RULE;
  ]                                                                                                                                                                                 
                                                                                                                                                                                    
  const DEFAULT_PASSES = STANDARD_PASSES

build_function(..., optimize = nothing) # fallback for no optimization
build_function(..., optimize = DEFAULT_PASSES)
...

Both should work mostly identically

[AayushSabharwal]

So we don't run into as much dynamic dispatch, could the global mapping be a NamedTuple and the passes in each level be a Tuple? SCP can also then add its own build targets, such as OptimizedJuliaTarget{T}() where T is the Symbol indicating the optimization level. Owning the build target would also enable SCP to add its own build_function kwargs without having to PR to Symbolics constantly.

[DhairyaLGandhi]

Ref SciML/ModelingToolkit.jl#4191