[tune][xegpu] infrastructure for tuning, applied to XeGPU matmul example#63
[tune][xegpu] infrastructure for tuning, applied to XeGPU matmul example#63
Conversation
| return parameters | ||
|
|
||
|
|
||
| matmul_param_db = { |
There was a problem hiding this comment.
Remove this dict definition, params are read from the json file now.
| PREFETCH_INST_DATA = [8, 16] | ||
| NB_WORKITEMS = 16 # workitems in subgroup | ||
| @KnobValue.ast_rewrite(in_exprs=True) | ||
| def checked_params_or_knobs( |
There was a problem hiding this comment.
Can we call this something like construct_search_space? I.e. for every entry params dict we either keep the given int value and check its validity or populate None values a knob with valid choices.
| NB_WORKITEMS = 16 # workitems in subgroup | ||
| @KnobValue.ast_rewrite(in_exprs=True) | ||
| def checked_params_or_knobs( | ||
| params: dict[str, int | None], layer_id="" |
There was a problem hiding this comment.
In this func we could rename layer_id -> knob_prefix because that's what we use it for. With the prefix every knob has an unique name; is this just a good practice/debugging aid, i.e. even if every layer had knobs with the same name everything would just work?
| if isinstance(sg_threads, smt_ext.SMTIntValue): | ||
| # NB: Constraint only enabled during tuning. |
There was a problem hiding this comment.
nit: I think you can remove the if now
| prefetch_a_k: int | KnobValue, | ||
| prefetch_b_k: int | KnobValue, | ||
| prefetch_b_n: int | KnobValue, | ||
| prefetch_nb: int, |
There was a problem hiding this comment.
prefetch_nb can be a knob too.
|
|
||
| def bundle_xegpu_to_binary(mod, stop_at_stage: str = "") -> ir.Module: | ||
| def bundle_xegpu_to_binary( | ||
| mod, stop_at_stage: str = "" |
| for tuneable_values in product( | ||
| *(tuneable.possibilities() for tuneable in tuneables) | ||
| ): | ||
| environment = dict(zip((tunable for tunable in tuneables), tuneable_values)) |
There was a problem hiding this comment.
| environment = dict(zip((tunable for tunable in tuneables), tuneable_values)) | |
| environment = dict(zip(tuneables, tuneable_values)) |
|
|
||
|
|
||
| class Node(ABC): | ||
| """Base class for `Node`s which can be evaluated w.r.t. an environment.""" |
There was a problem hiding this comment.
For better clarity a high level docstring would help: We are tracing the IR to construct knob dependency tree, "Node" represents a node in that tree; "environment" means the (tunable, value) map.
| parser.add_argument( | ||
| "--mode", | ||
| choices=["enumerate"], | ||
| default="enumerate", | ||
| help="Mode of operation", | ||
| ) |
There was a problem hiding this comment.
nit: add --mode once we have more than one mode
2 participants
Adds ability to extract search spaces from schedules with knobs.
Adds transform interpreter semantics to
constrain_paramsop, which interprets the smt ops in its region so that it is possible to check constraints on int params as well as calculate new int params at interpreter time.Modifies the schedule for the XeGPU matmul (& mlp) example to embed its tuning problem.
lighthouse.tuneimplements:trace-ing of transform schedules and the SMT regions they include:Nodes, eachevaluate-able w.r.t. an environment, with leafs such asConstant, which represents just a constant in IR, and evals to a constantint, andKnob, the representative of atransform.tune.knob, which takes its value from the env while knowing what it's possible values are,while
Applydepends on otherNodes as it modelsValues produced by ops dependent on other valuesPredicatemodels which condition/constraint needs to be true, as a boolean-valued function onNodes, for execution to proceed passed a particular op.rewrite-ing of transform schedules, solely through setting theselectedattr onknobops and theselected_regionattr onalternativesops.enumerate-ing all valid assignments forknobandalternativestuneables.__main__-ing to take a .mlir schedule and derive all valid knob configurations and output the corresponding concrete/transform-interpreter interpretable schedules.Inside
lighthouse.dialects, add (extension) dialects:smt_ext: A wrapper forir.Values of!smt.inttype so we can support python operations on them (e.g. addition) also with integers.transform_tune_ext: A wrapper forir.Values produced directly bytransform.tune.knobops so we can do python operations on them, in particular add constraints, and a camel_caseknob() -> KnobValuehelpertransform_smt_ext: so we we can have a version oftransform.smt.constrain_paramswhich has transform-interpreter semantics: by tracing the body, containing smt ops, we get function we can applied to thetransform.paramwhich were arguments toconstrain_params. That is, this version ofconstrain_paramshas a properTransformOpInterfaceimplementation.