[Tune] Constraint extraction, model finding, and schedule rewriting

[rolfmorel]

For a non-trivial example, run python examples/xegpu_matmul/matmul.py --dump-schedule --non-det for

module attributes {transform.with_named_sequence} {                                                                                                                                  
  transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {                                                                                        
    %0 = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op                                                                           
    %1 = transform.get_parent_op %0 {deduplicate, op_name = "builtin.module"} : (!transform.any_op) -> !transform.any_op                                                             
    %2 = transform.tune.knob<"wg_d0"> options = {lb = 128 : i64, step = 32 : i64, ub = 256 : i64} -> !transform.any_param                                                            
    %3 = transform.tune.knob<"wg_d1"> options = {lb = 128 : i64, step = 32 : i64, ub = 256 : i64} -> !transform.any_param                                                            
    %4 = transform.tune.knob<"sg_d0"> options = {lb = 16 : i64, step = 8 : i64, ub = 32 : i64} -> !transform.any_param                                                               
    %5 = transform.tune.knob<"sg_d1"> options = {lb = 16 : i64, step = 8 : i64, ub = 32 : i64} -> !transform.any_param                                                               
    %6:3 = transform.smt.constrain_params(%2, %3, %4, %5) : (!transform.any_param, !transform.any_param, !transform.any_param, !transform.any_param) -> (!transform.any_param, !trans
form.any_param, !transform.any_param) {                                                                                                                                              
    ^bb0(%arg1: !smt.int, %arg2: !smt.int, %arg3: !smt.int, %arg4: !smt.int):                                                                                                        
      %c0 = smt.int.constant 0                                                                                                                                                       
      %58 = smt.int.mod %arg1, %arg3
...

then feed that to ... | python -m lighthouse.tune - --print-smtlib -n 100:

////////////////////////////////////////////////////////////////////////////////                                                                                                     
// SMT-LIB constraints                                                                                                                                                               
////////////////////////////////////////////////////////////////////////////////                                                                                                     
(declare-fun wg_d0 () Int)                                                                                                                                                           
(declare-fun wg_d1 () Int)                                                                                                                                                           
(declare-fun sg_d0 () Int)                                                                                                                                                           
(declare-fun sg_d1 () Int)                                                                                                                                                           
(declare-fun cp_bbarg4 () Int)
...
(assert (and (>= wg_d0 128) (<= wg_d0 256) (= (mod wg_d0 32) 0)))
(assert (and (>= wg_d1 128) (<= wg_d1 256) (= (mod wg_d1 32) 0)))
(assert (and (>= sg_d0 16) (<= sg_d0 32) (= (mod sg_d0 8) 0)))
(assert (and (>= sg_d1 16) (<= sg_d1 32) (= (mod sg_d1 8) 0)))
(assert (and (= cp_bbarg1 wg_d0)
     (= cp_bbarg2 wg_d1)
     (= cp_bbarg3 sg_d0)
     (= cp_bbarg4 sg_d1)))
...
////////////////////////////////////////////////////////////////////////////////
// Determinized schedule #1
////////////////////////////////////////////////////////////////////////////////
module attributes {transform.with_named_sequence} {
  transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
    %0 = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op
    %1 = transform.get_parent_op %0 {deduplicate, op_name = "builtin.module"} : (!transform.any_op) -> !transform.any_op
    %2 = transform.param.constant 160 : i64 -> !transform.param<i64>
    %3 = transform.param.constant 192 : i64 -> !transform.param<i64>
    %4 = transform.param.constant 16 : i64 -> !transform.param<i64>
    %5 = transform.param.constant 24 : i64 -> !transform.param<i64>
...
////////////////////////////////////////////////////////////////////////////////                                                                                                     
// Determinized schedule #100                                                                                                                                                        
////////////////////////////////////////////////////////////////////////////////                                                                                                     
module attributes {transform.with_named_sequence} {                                                                                                                                  
  transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {                                                                                        
    %0 = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op                                                                           
    %1 = transform.get_parent_op %0 {deduplicate, op_name = "builtin.module"} : (!transform.any_op) -> !transform.any_op                                                             
    %2 = transform.param.constant 256 : i64 -> !transform.param<i64>                                                                                                                 
    %3 = transform.param.constant 256 : i64 -> !transform.param<i64>                                                                                                                 
    %4 = transform.param.constant 16 : i64 -> !transform.param<i64>                                                                                                                  
    %5 = transform.param.constant 32 : i64 -> !transform.param<i64>
...

Co-authored-by: Tuomas Kärnä [email protected]

[tkarna]

It would be good if we could still dump both the payload and schedule simultaneously if so desired.

[tkarna]

Maybe a more descriptive name e.g. "Unselected"

[tkarna]

Apart from a couple of minor comments this looks good time!